7f1aee6dc8
* enable N rules for naming conventions * fix pyi * address comments * remove unneeded code
136 lines
3.7 KiB
Python
136 lines
3.7 KiB
Python
"""Checks the size of a specific directory and uploads result to Posthog."""
|
|
|
|
import argparse
|
|
import os
|
|
from pathlib import Path
|
|
|
|
from utils import get_directory_size, get_python_version, send_data_to_posthog
|
|
|
|
|
|
def get_package_size(venv_path: Path, os_name):
|
|
"""Get the size of a specified package.
|
|
|
|
Args:
|
|
venv_path: The path to the venv.
|
|
os_name: Name of os.
|
|
|
|
Returns:
|
|
The total size of the package in bytes.
|
|
|
|
Raises:
|
|
ValueError: when venv does not exist or python version is None.
|
|
"""
|
|
python_version = get_python_version(venv_path, os_name)
|
|
print("Python version:", python_version)
|
|
if python_version is None:
|
|
raise ValueError("Error: Failed to determine Python version.")
|
|
|
|
is_windows = "windows" in os_name
|
|
|
|
package_dir: Path = (
|
|
venv_path / "lib" / f"python{python_version}" / "site-packages"
|
|
if not is_windows
|
|
else venv_path / "Lib" / "site-packages"
|
|
)
|
|
if not package_dir.exists():
|
|
raise ValueError(
|
|
"Error: Virtual environment does not exist or is not activated."
|
|
)
|
|
|
|
total_size = get_directory_size(package_dir)
|
|
return total_size
|
|
|
|
|
|
def insert_benchmarking_data(
|
|
os_type_version: str,
|
|
python_version: str,
|
|
commit_sha: str,
|
|
pr_title: str,
|
|
branch_name: str,
|
|
pr_id: str,
|
|
path: str,
|
|
):
|
|
"""Insert the benchmarking data into PostHog.
|
|
|
|
Args:
|
|
os_type_version: The OS type and version to insert.
|
|
python_version: The Python version to insert.
|
|
commit_sha: The commit SHA to insert.
|
|
pr_title: The PR title to insert.
|
|
branch_name: The name of the branch.
|
|
pr_id: The id of the PR.
|
|
path: The path to the dir or file to check size.
|
|
"""
|
|
if "./dist" in path:
|
|
size = get_directory_size(Path(path))
|
|
else:
|
|
size = get_package_size(Path(path), os_type_version)
|
|
|
|
# Prepare the event data
|
|
properties = {
|
|
"path": path,
|
|
"os": os_type_version,
|
|
"python_version": python_version,
|
|
"distinct_id": commit_sha,
|
|
"pr_title": pr_title,
|
|
"branch_name": branch_name,
|
|
"pr_id": pr_id,
|
|
"size_mb": round(
|
|
size / (1024 * 1024), 3
|
|
), # save size in MB and round to 3 places
|
|
}
|
|
|
|
send_data_to_posthog("package_size", properties)
|
|
|
|
|
|
def main():
|
|
"""Runs the benchmarks and inserts the results."""
|
|
parser = argparse.ArgumentParser(description="Run benchmarks and process results.")
|
|
parser.add_argument(
|
|
"--os", help="The OS type and version to insert into the database."
|
|
)
|
|
parser.add_argument(
|
|
"--python-version", help="The Python version to insert into the database."
|
|
)
|
|
parser.add_argument(
|
|
"--commit-sha", help="The commit SHA to insert into the database."
|
|
)
|
|
parser.add_argument(
|
|
"--pr-title",
|
|
help="The PR title to insert into the database.",
|
|
)
|
|
parser.add_argument(
|
|
"--branch-name",
|
|
help="The current branch",
|
|
required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--pr-id",
|
|
help="The pr id",
|
|
required=True,
|
|
)
|
|
parser.add_argument(
|
|
"--path",
|
|
help="The path to the vnenv.",
|
|
required=True,
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Get the PR title from env or the args. For the PR merge or push event, there is no PR title, leaving it empty.
|
|
pr_title = args.pr_title or os.getenv("PR_TITLE", "")
|
|
|
|
# Insert the data into the database
|
|
insert_benchmarking_data(
|
|
os_type_version=args.os,
|
|
python_version=args.python_version,
|
|
commit_sha=args.commit_sha,
|
|
pr_title=pr_title,
|
|
branch_name=args.branch_name,
|
|
pr_id=args.pr_id,
|
|
path=args.path,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|