diff --git a/core/config.py b/core/config.py index 8a940e7a..6cfe3405 100644 --- a/core/config.py +++ b/core/config.py @@ -12,6 +12,7 @@ class PackageManager(Enum): CRATES = "crates" HOMEBREW = "homebrew" + PYPI = "pypi" TEST = env_vars("TEST", "false") @@ -20,6 +21,7 @@ class PackageManager(Enum): SOURCES = { PackageManager.CRATES: "https://static.crates.io/db-dump.tar.gz", PackageManager.HOMEBREW: "https://github.com/Homebrew/homebrew-core/tree/master/Formula", # noqa + PackageManager.PYPI: "https://pypi.org/simple/", # Base URL for PyPI's Simple API } # The three configuration values URLTypes, DependencyTypes, and UserTypes will query the diff --git a/core/db.py b/core/db.py index 33b649f8..6cf04f9e 100644 --- a/core/db.py +++ b/core/db.py @@ -128,7 +128,7 @@ def update_caches( ): if update_packages: self._update_cache( - self.package_cache, Package, "import_id", "id", items, "crate_id" + self.package_cache, Package, "import_id", "id", items, "import_id" ) if update_users: self._update_cache( @@ -136,7 +136,7 @@ def update_caches( ) if update_versions: self._update_cache( - self.version_cache, Version, "import_id", "id", items, "version_id" + self.version_cache, Version, "import_id", "id", items, "import_id" ) if update_licenses: self._update_cache( @@ -159,10 +159,15 @@ def insert_versions(self, version_generator: Iterable[dict[str, str]]): self._insert_batch(Version, versions) def _process_version(self, item: Dict[str, str]): - package_id = self.package_cache.get(item["crate_id"]) - if not package_id: - self.logger.warn(f"package {item['crate_id']} not found") - return None + # FIXME: this is a hack, the import_id of a version shouldn't be the same as the package's import_id + # but the original logic here tries to use the package's import_id as the version's import_id + # this is a temporary fix (for the implementation of pypi) + # to fix this completely, we need to update the logic of the crates transformer too + if not (package_id := item.get("package_id")): + package_id = self.package_cache.get(item["import_id"]) + if not package_id: + self.logger.warn(f"package {item['import_id']} not found") + return None license_id = self.license_cache.get(item["license"]) if not license_id: @@ -201,11 +206,31 @@ def insert_dependencies(self, dependency_generator: Iterable[dict[str, str]]): self._insert_batch(DependsOn, dependencies) def _process_depends_on(self, item: Dict[str, str]): - return DependsOn( - version_id=self.version_cache[item["version_id"]], - dependency_id=self.package_cache[item["crate_id"]], - semver_range=item["semver_range"], - ).to_dict() + version_id = self.version_cache.get(item["version_id"]) + + # in case the version cannot be found from the cache + if not version_id: + # we need to fetch from the database + version = self.select_version_by_import_id(item["version_id"]) + if not version: + self.logger.warn(f"version {item['version_id']} not found") + return None + version_id = version.id + self.version_cache[item["version_id"]] = version_id + + # Create base dependency object + depends_on = { + "version_id": self.version_cache[item["version_id"]], + "dependency_id": self.package_cache[item["import_id"]], + "semver_range": item["semver_range"] + } + + # Add dependency_type_id if provided + depends_on.update( + {"dependency_type_id": item["dependency_type_id"]} if "dependency_type_id" in item else {} + ) + + return DependsOn(**depends_on).to_dict() def insert_users(self, user_generator: Iterable[dict[str, str]], source_id: UUID): def process_user(item: Dict[str, str]): @@ -245,13 +270,13 @@ def _process_user_package(self, item: Dict[str, str]): self.logger.warn(f"user {item['owner_id']} not found") return None - if item["crate_id"] not in self.package_cache: - self.logger.warn(f"package {item['crate_id']} not found") + if item["import_id"] not in self.package_cache: + self.logger.warn(f"package {item['import_id']} not found") return None return UserPackage( user_id=self.user_cache[item["owner_id"]], - package_id=self.package_cache[item["crate_id"]], + package_id=self.package_cache[item["import_id"]], ).to_dict() def insert_user_versions( @@ -476,6 +501,18 @@ def select_version_by_import_id(self, import_id: str) -> Version | None: result = session.query(Version).filter_by(import_id=import_id).first() if result: return result + + def select_latest_version_by_import_id(self, import_id: str) -> Version | None: + with self.session() as session: + # First get the package + package = session.query(Package).filter_by(import_id=import_id).first() + if not package: + return None + + # Then get the latest version for this package + result = session.query(Version).filter_by(package_id=package.id).order_by(Version.version.desc()).first() + if result: + return result def select_package_manager_name_by_id(self, id: UUID) -> str | None: with self.session() as session: diff --git a/core/models/__init__.py b/core/models/__init__.py index 1ef17e4f..ea0ec342 100644 --- a/core/models/__init__.py +++ b/core/models/__init__.py @@ -189,13 +189,18 @@ class DependsOn(Base): dependency_type: Mapped["DependsOnType"] = relationship() def to_dict(self): - return { + result = { "version_id": self.version_id, "dependency_id": self.dependency_id, - # "dependency_type_id": self.dependency_type_id, "semver_range": self.semver_range, } + # if dependency_type_id is provided, include it + if self.dependency_type_id is not None: + result["dependency_type_id"] = self.dependency_type_id + + return result + class DependsOnType(Base): __tablename__ = "depends_on_types" diff --git a/docker-compose.yml b/docker-compose.yml index 05e3debe..b8e5a693 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -69,6 +69,26 @@ services: alembic: condition: service_completed_successfully + pypi: + build: + context: . + dockerfile: ./package_managers/pypi/Dockerfile + environment: + - CHAI_DATABASE_URL=postgresql://postgres:s3cr3t@db:5432/chai + - NO_CACHE=${NO_CACHE:-false} + - PYTHONPATH=/ + - DEBUG=${DEBUG:-false} + - TEST=${TEST:-false} + - FETCH=${FETCH:-true} + - FREQUENCY=${FREQUENCY:-24} + volumes: + - ./data/pypi:/data/pypi + depends_on: + db: + condition: service_healthy + alembic: + condition: service_completed_successfully + api: build: context: ./api diff --git a/package_managers/crates/transformer.py b/package_managers/crates/transformer.py index d2b79da3..001e7828 100644 --- a/package_managers/crates/transformer.py +++ b/package_managers/crates/transformer.py @@ -92,7 +92,7 @@ def dependencies(self) -> Generator[Dict[str, str], None, None]: yield { "version_id": start_id, - "crate_id": end_id, + "import_id": end_id, "semver_range": req, "dependency_type": dependency_type, } @@ -130,7 +130,7 @@ def user_packages(self) -> Generator[Dict[str, str], None, None]: owner_id = row["owner_id"] yield { - "crate_id": crate_id, + "import_id": crate_id, "owner_id": owner_id, } diff --git a/package_managers/pypi/Dockerfile b/package_managers/pypi/Dockerfile new file mode 100644 index 00000000..7da0c000 --- /dev/null +++ b/package_managers/pypi/Dockerfile @@ -0,0 +1,10 @@ +FROM python:3.11-slim + +WORKDIR /app + +COPY package_managers/pypi/requirements.txt . +RUN pip install -r requirements.txt + +COPY . . + +CMD ["python", "-m", "package_managers.pypi.main"] diff --git a/package_managers/pypi/README.md b/package_managers/pypi/README.md new file mode 100644 index 00000000..356281d7 --- /dev/null +++ b/package_managers/pypi/README.md @@ -0,0 +1,31 @@ +# PyPI + +The PyPI service processes package data from the Python Package Index (PyPI) and transforms it into CHAI's normalized format. It uses PyPI's JSON data dumps for efficient bulk processing. + +## Getting Started + +To run the PyPI service, use the following commands: + +```bash +docker compose build pypi +docker compose run pypi +``` + +## Execution Steps + +The PyPI loader follows these steps: + +1. Initialization: Sets up configuration and database connection +2. Fetching: Downloads the latest PyPI JSON data dump if `FETCH` is true +3. Transformation: Converts PyPI's JSON format into CHAI's schema +4. Loading: Inserts transformed data into the database: + - Packages + - Users + - User Packages + - URLs + - Package URLs + - Versions + - Dependencies +5. Cleanup: Removes temporary files if `NO_CACHE` is true + +The main execution logic is in the `run_pipeline` function in `main.py`. diff --git a/package_managers/pypi/fetcher.py b/package_managers/pypi/fetcher.py new file mode 100644 index 00000000..1183b026 --- /dev/null +++ b/package_managers/pypi/fetcher.py @@ -0,0 +1,244 @@ +from concurrent.futures import ThreadPoolExecutor, as_completed +from typing import Generator, List, Optional, Any +from html.parser import HTMLParser +from dataclasses import dataclass +from urllib.parse import urljoin +from datetime import datetime +import multiprocessing +import json +import time +import os + +import requests + +from core.config import Config +from core.fetcher import Data, Fetcher +from core.logger import Logger + + +class SimpleIndexParser(HTMLParser): + """Parser for PyPI's simple index page to extract package names.""" + def __init__(self): + super().__init__() + self.packages = [] + + def handle_starttag(self, tag: str, attrs: List[tuple]): + if tag == "a": + for attr, value in attrs: + if attr == "href": + # Package names are the href values + self.packages.append(value) + + +class PyPIFetcher(Fetcher): + """Custom fetcher for PyPI that handles both simple index and JSON API.""" + + def __init__(self, name: str, config: Config): + super().__init__(name, config) + self.session = requests.Session() + self.base_url = "https://pypi.org" + self.rate_limit_delay = 0.1 # seconds between requests + self.batch_size = 100 # packages per batch + self.max_workers = multiprocessing.cpu_count() * 4 # Number of threads for parallel downloads + self.data_dir = "/data/pypi" # We will mount to ./data/pypi on the host + self.process_file = os.path.join(self.data_dir, "progress.json") # to store the progress of the fetch + self.packages_file = os.path.join(self.data_dir, "packages.txt") # a list of all package names + + def _save_process(self, batch_num: int, downloaded: int, fetched: int, total: int) -> None: + """Save current process to progress.json""" + process = { + "batch_num": batch_num, + "downloaded": downloaded, # successful downloads + "fetched": fetched, # total attempted downloads + "total": total, + "timestamp": datetime.now().isoformat() + } + try: + with open(self.process_file, 'w') as f: + json.dump(process, f, indent=2) + except (IOError, OSError) as e: + self.logger.error(f"Failed to open process file for writing: {e}") + except json.JSONEncodeError as e: + self.logger.error(f"Failed to serialize process data to JSON: {e}") + + def _load_process(self) -> tuple[int, int, int, int]: + """Load process from progress.json if exists""" + try: + with open(self.process_file) as f: + process = json.load(f) + return ( + process["batch_num"], + process.get("downloaded", 0), + process.get("fetched", 0), + process["total"] + ) + except (FileNotFoundError, json.JSONDecodeError) as e: + self.logger.error(f"Failed to load process file: {e}") + return 0, 0, 0, 0 + + def _save_package_list(self, packages: List[str]): + """Save package list to packages.txt""" + try: + with open(self.packages_file, 'w') as f: + f.writelines(f"{package}\n" for package in packages) + except (IOError, OSError) as e: + self.logger.error(f"Failed to open package list file for writing: {e}") + + def _load_package_list(self) -> List[str]: + """Load package list from packages.txt if exists""" + try: + with open(self.packages_file, 'r') as f: + return [line.strip() for line in f if line.strip()] + except FileNotFoundError: + return [] + except (IOError, OSError) as e: + self.logger.error(f"Failed to load package list file: {e}") + return [] + + def _get_package_list(self) -> List[str]: + """Fetch list of all packages from PyPI simple index.""" + url = urljoin(self.base_url, "simple") # PyPI simple index + try: + response = self.session.get(url) + response.raise_for_status() + except requests.RequestException as e: + self.logger.error(f"Failed to fetch package list from {url}") + self.logger.error(f"Status code: {response.status_code if 'response' in locals() else 'N/A'}") + self.logger.error(f"Response body: {response.text if 'response' in locals() else 'N/A'}") + self.logger.error(f"Error: {str(e)}") + return [] + + parser = SimpleIndexParser() + parser.feed(response.text) + return parser.packages + + def _get_package_data(self, package_name: str) -> dict[str, Any] | None: + """Fetch JSON data for a specific package.""" + # Remove any /simple/ prefix if present + package_name = package_name.replace('/simple/', '') + # Remove any trailing slash + package_name = package_name.rstrip('/') + + base_api_url = urljoin(self.base_url, "pypi/") + package_url = urljoin(base_api_url, f"{package_name}/") + url = urljoin(package_url, "json") + + try: + response = self.session.get(url) + response.raise_for_status() + return response.json() + except requests.RequestException as e: + self.logger.error(f"Error fetching {package_name}: {e}") + return None + + def _download_batch(self, packages: List[str]) -> List[dict]: + """Download a batch of packages in parallel.""" + results = [] + with ThreadPoolExecutor(max_workers=self.max_workers) as executor: + future_to_package = { + executor.submit(self._get_package_data, package): package + for package in packages + } + for future in as_completed(future_to_package): + package = future_to_package[future] + try: + data = future.result() + if data: + results.append(data) + except Exception as e: + self.logger.error(f"Error downloading {package}: {e}") + return results + + def fetch(self) -> Generator[Data, None, None]: + """ + Fetch package data from PyPI: + 1. Check if we need to resume or start fresh + 2. Get/load list of all packages + 3. Split packages into batches + 4. Download each batch in parallel + 5. Save raw data to /data/pypi/[batch_num].json in container + 6. Track progress in progress.json + """ + # Create data directory if needed + try: + os.makedirs(self.data_dir, exist_ok=True) + except (IOError, OSError) as e: + self.logger.error(f"Failed to create data directory: {e}") + return + + # Check if we need to resume + current_batch, downloaded, fetched, total = self._load_process() + + # If everything is fetched, don't do anything + # TODO: implement a re-download logic + if fetched == total and total > 0: + self.logger.log(f"All packages already fetched ({downloaded}/{fetched}/{total})") + return + + # Get package list + if current_batch == 0: # Only get new list if fresh start + self.logger.log("Starting fresh download") + packages = self._get_package_list() + self._save_package_list(packages) + downloaded = 0 + fetched = 0 + else: + # Resume from previous run + self.logger.log(f"Resuming from batch {current_batch} ({downloaded}/{fetched}/{total} packages)") + packages = self._load_package_list() + if not packages: + self.logger.log("No saved package list found, starting fresh") + packages = self._get_package_list() + self._save_package_list(packages) + downloaded = 0 + fetched = 0 + + total_packages = len(packages) + self.logger.log(f"Found {total_packages} packages") + + # Skip already downloaded batches + start_idx = current_batch * self.batch_size + packages = packages[start_idx:] + + # Process remaining packages in batches + for i in range(0, len(packages), self.batch_size): + batch = packages[i:i + self.batch_size] + batch_num = (start_idx + i)//self.batch_size + 1 + total_batches = (total_packages + self.batch_size - 1)//self.batch_size + self.logger.log(f"Downloading batch {batch_num}/{total_batches}") + + # Download batch in parallel + results = self._download_batch(batch) + + # Update counters + fetched = start_idx + i + len(batch) # Count all attempted packages + downloaded += len(results) if results else 0 # Count only successful downloads + + if results: + # Save batch to JSON file + file_name = f"{batch_num}.json" + file_path = os.path.join(self.data_dir, file_name) + + try: + with open(file_path, 'w') as f: + json.dump(results, f) + except (IOError, OSError) as e: + self.logger.error(f"Failed to open batch file for writing: {e}") + except json.JSONEncodeError as e: + self.logger.error(f"Failed to serialize batch data to JSON: {e}") + + # Update progress + self._save_process(batch_num, downloaded, fetched, total_packages) + + # Yield Data object for tracking + yield Data( + file_path=self.data_dir, + file_name=file_name, + content=json.dumps(results).encode('utf-8') + ) + else: + # Still save progress even if no results + self._save_process(batch_num, downloaded, fetched, total_packages) + + # Rate limiting between batches + time.sleep(self.rate_limit_delay) diff --git a/package_managers/pypi/main.py b/package_managers/pypi/main.py new file mode 100644 index 00000000..c2bf7366 --- /dev/null +++ b/package_managers/pypi/main.py @@ -0,0 +1,101 @@ +import time +import json +from typing import Generator + +from core.config import Config, PackageManager +from core.db import DB +from core.logger import Logger +from core.scheduler import Scheduler +from package_managers.pypi.fetcher import PyPIFetcher, Data +from package_managers.pypi.transformer import PyPITransformer + +logger = Logger("pypi_orchestrator") + + +def fetch(config: Config) -> Generator[Data, None, None]: + logger.log("šŸ”„ Starting PyPI data fetching process...") + fetcher = PyPIFetcher("pypi", config) + + if config.exec_config.fetch: + logger.log("šŸ“„ Fetching new data from PyPI...") + yield from fetcher.fetch() + else: + logger.log("ā„¹ļø Skipping fetch (FETCH=false)") + return # Empty generator + + +def run_pipeline(db: DB, config: Config) -> None: + """Run the PyPI pipeline.""" + logger.log("\nšŸš€ Starting PyPI pipeline...") + + # Download all the package data from PyPI, and process them after that + # As of Jan 2025, PyPI has ~600k packages + logger.log("\nšŸ”„ Starting download pipeline...") + for data in fetch(config): + logger.log(f"Saved batch to {data.file_name}") + + # Create transformer + transformer = PyPITransformer( + url_types=config.url_types, + user_types=config.user_types, + pm_config=config.pm_config, + db=db + ) + + # Step 1: Insert packages and related data + logger.log("\nšŸ“¦ Inserting packages...") + db.insert_packages( + transformer.packages(), + config.pm_config.pm_id, + PackageManager.PYPI.value, + ) + + # Step 2: Insert licenses + # This is handled during package insertion as we create licenses on demand + + # Step 3: Skip user-related operations as we can't get GitHub info from PyPI + + # Step 4: Insert URLs and package URLs + logger.log("\nšŸ”— Inserting URLs...") + db.insert_urls(transformer.urls()) + + logger.log("\nšŸ”— Linking URLs and Packages...") + db.insert_package_urls(transformer.package_urls()) + + # Step 5: Insert versions + logger.log("\nšŸ“ Inserting versions...") + db.insert_versions(transformer.versions()) + + # Step 6: Insert dependencies (after all packages are in) + logger.log("\nšŸ”„ Inserting dependencies...") + db.insert_dependencies(transformer.dependencies()) + + # Record load history + db.insert_load_history(config.pm_config.pm_id) + logger.log("āœ… PyPI data loading completed successfully") + + +def main(): + logger.log("\nšŸ“¦ Initializing PyPI Package Manager...") + db = DB() + config = Config(PackageManager.PYPI, db) + logger.debug(f"Configuration: {config}") + + scheduler = Scheduler("pypi") + scheduler.start(run_pipeline, db, config) + + # run immediately + scheduler.run_now(run_pipeline, db, config) + + # keep the main thread alive so we can terminate the program with Ctrl+C + try: + while True: + time.sleep(1) + except KeyboardInterrupt: + logger.log("\nāš ļø Received interrupt signal, shutting down...") + scheduler.stop() + logger.log("āœ… Shutdown complete") + + +if __name__ == "__main__": + main() diff --git a/package_managers/pypi/requirements.txt b/package_managers/pypi/requirements.txt new file mode 100644 index 00000000..94c044a1 --- /dev/null +++ b/package_managers/pypi/requirements.txt @@ -0,0 +1,7 @@ +alembic>=1.13.2 +psycopg2-binary>=2.9.9 +requests>=2.32.3 +schedule>=1.2.1 +sqlalchemy>=2.0.23 +pyyaml>=6.0.2 +beautifulsoup4>=4.12.2 # For parsing HTML from PyPI's simple index diff --git a/package_managers/pypi/structs.py b/package_managers/pypi/structs.py new file mode 100644 index 00000000..e2e7a667 --- /dev/null +++ b/package_managers/pypi/structs.py @@ -0,0 +1,10 @@ +from enum import IntEnum + +class DependencyType(IntEnum): + """Enum for PyPI dependency types.""" + RUNTIME = 1 # Default runtime dependency + DEV = 2 # Development dependency + OPTIONAL = 3 # Optional dependency (e.g., extras) + + def __str__(self): + return self.name.lower() diff --git a/package_managers/pypi/transformer.py b/package_managers/pypi/transformer.py new file mode 100644 index 00000000..7568d025 --- /dev/null +++ b/package_managers/pypi/transformer.py @@ -0,0 +1,471 @@ +import json +import glob +import os +from typing import Dict, Generator, List, Any +from pathlib import Path + +from core.config import URLTypes, UserTypes, PMConf +from core.transformer import Transformer +from core.utils import safe_int +from core.db import DB +from core.logger import Logger +from core.models import URL, PackageURL +from package_managers.pypi.structs import DependencyType + + +class PyPITransformer(Transformer): + """Transform PyPI package data into CHAI's package format.""" + + def __init__(self, url_types: URLTypes, user_types: UserTypes, pm_config: PMConf, db: DB, data_dir: str = "/data/pypi"): + """Initialize PyPI transformer.""" + super().__init__("pypi") + self.url_types = url_types + self.user_types = user_types + self.pm_config = pm_config + self.db = db + self.data_dir = data_dir + self.logger = Logger("pypi_transformer") + + # Load progress to know how many batches to process + self.progress_file = os.path.join(data_dir, "progress.json") + try: + with open(self.progress_file, 'r') as f: + self.progress = json.load(f) + self.total_batches = self.progress["batch_num"] + except (FileNotFoundError, json.JSONDecodeError) as e: + self.logger.error(f"Error loading progress.json: {e}") + self.progress = None + # Count JSON files in directory + json_files = glob.glob(os.path.join(data_dir, "[0-9]*.json")) + self.total_batches = len(json_files) + + def _read_batch_files(self) -> Generator[Dict, None, None]: + """ + Process package data from JSON files in data_dir. + Files are named 1.json, 2.json, etc. + Yields each package's data. + """ + total_packages = 0 + self.logger.log(f"Processing {self.total_batches} batch files") + + for batch_num in range(1, self.total_batches + 1): + file_path = os.path.join(self.data_dir, f"{batch_num}.json") + try: + with open(file_path) as f: + packages = json.load(f) + self.logger.log(f"Processing batch {batch_num} with {len(packages)} packages") + for package in packages: + yield package + total_packages += 1 + except (FileNotFoundError, json.JSONDecodeError) as e: + self.logger.error(f"Error processing {file_path}: {e}") + continue + + self.logger.log(f"Total packages processed by transformer: {total_packages}") + + def packages(self) -> Generator[Dict[str, Any], None, None]: + """Transform PyPI package data into CHAI's package format.""" + for package_data in self._read_batch_files(): + try: + info = package_data.get("info") + try: + assert info, "Package info is missing" + except AssertionError as e: + self.logger.warn(str(e)) + continue + + name = info.get("name") + if not name: + continue + + package_data = { + "derived_id": f"pypi/{name}", + "name": name, + "package_manager_id": self.pm_config.pm_id, + "import_id": f"pypi/{name}", + "readme": info.get("description", "") + } + + # Update package cache immediately + self.db.package_cache[f"pypi/{name}"] = package_data + + yield package_data + except Exception as e: + self.logger.error(f"Error processing package: {e}") + continue + + def urls(self) -> Generator[Dict[str, str], None, None]: + """Transform PyPI URLs into CHAI's URL format.""" + + for package_data in self._read_batch_files(): + try: + info = package_data.get("info", {}) + + # Extract URLs from various fields + urls = [] + + # Homepage URL + homepage = info.get("home_page") + if homepage and isinstance(homepage, str): + homepage = homepage.strip() + if homepage: + urls.append((homepage, self.url_types.homepage)) + + # Documentation URL + docs_url = info.get("docs_url") + if docs_url and isinstance(docs_url, str): + docs_url = docs_url.strip() + if docs_url: + urls.append((docs_url, self.url_types.documentation)) + + # Project URLs + project_urls = info.get("project_urls", {}) + if project_urls and isinstance(project_urls, dict): + for label, url in project_urls.items(): + if not url or not isinstance(url, str): + continue + + url = url.strip() + if not url: + continue + + # Map common project URL labels to URL types + url_type = None + label_lower = label.lower() + if "source" in label_lower or "repository" in label_lower or "github" in label_lower: + url_type = self.url_types.repository + elif "doc" in label_lower: + url_type = self.url_types.documentation + elif "home" in label_lower: + url_type = self.url_types.homepage + + if url_type: + urls.append((url, url_type)) + + # yield URLs + for url, url_type in urls: + yield { + "url": url, + "url_type_id": url_type + } + + except Exception as e: + self.logger.error(f"Error processing URLs: {e}") + continue + + def package_urls(self) -> Generator[Dict[str, str], None, None]: + """Transform PyPI package URLs into CHAI's package URL format.""" + for package_data in self._read_batch_files(): + try: + info = package_data.get("info", {}) + name = info.get("name") + if not name: + continue + + import_id = f"pypi/{name}" + + # Homepage URL + homepage = info.get("home_page") + if homepage and isinstance(homepage, str): + homepage = homepage.strip() + if homepage: + yield { + "import_id": import_id, + "url": homepage, + "url_type_id": self.url_types.homepage + } + + # Documentation URL + docs_url = info.get("docs_url") + if docs_url and isinstance(docs_url, str): + docs_url = docs_url.strip() + if docs_url: + yield { + "import_id": import_id, + "url": docs_url, + "url_type_id": self.url_types.documentation + } + + # Project URLs + project_urls = info.get("project_urls", {}) + if project_urls and isinstance(project_urls, dict): + for label, url in project_urls.items(): + if not url or not isinstance(url, str): + continue + + url = url.strip() + if not url: + continue + + # Map common project URL labels to URL types + url_type = None + label_lower = label.lower() + if "source" in label_lower or "repository" in label_lower or "github" in label_lower: + url_type = self.url_types.repository + elif "doc" in label_lower: + url_type = self.url_types.documentation + elif "home" in label_lower: + url_type = self.url_types.homepage + + if url_type: + yield { + "import_id": import_id, + "url": url, + "url_type_id": url_type + } + + except Exception as e: + self.logger.error(f"Error processing package URLs: {e}") + continue + + def versions(self) -> Generator[Dict[str, Any], None, None]: + """Transform PyPI versions into CHAI's format.""" + + for package_data in self._read_batch_files(): + try: + info = package_data.get("info") + try: + assert info, "Package info is missing" + except AssertionError as e: + self.logger.warn(str(e)) + continue + + name = info.get("name") + pypi_id = f"pypi/{name}" + + package = self.db.select_package_by_import_id(pypi_id) + if not package: + self.logger.warn(f"Package {pypi_id} not found in database") + continue + + # We are reading from .releases, the object key is the version number + # We need to loop through all the versions from the .releases object + for version, releases in package_data.get("releases", {}).items(): + # Check if the release array is empty + if len(releases) == 0: + continue + + # As there can be multiple releases for a version, we will take the first one + release = releases[0] + + # Get license + license_name = info.get("license") or "Unknown" + + # Some license has set to a full paragraph + # We need to check if the length is greater than 50, and try to extract the name with a comma + # And get the first part of that + if len(license_name) > 50: + license_name = license_name.split(",")[0] + + # If it still has more than 50 characters, just skip it + if len(license_name) > 50: + continue + + # Get downloads + downloads = release.get("downloads", -1) + + # Get size + size = release.get("size", 0) + + # Get published_at + published_at = release.get("upload_time_iso_8601", "") + + # Get checksum + checksum = release.get("digests", {}).get("sha256") + + # Yield data + yield { + "version": version, + "package_id": package.id, + "import_id": f"{pypi_id}-{version}", + "size": size, + "published_at": published_at, + "license": license_name, + "downloads": downloads, + "checksum": checksum, + } + + except Exception as e: + self.logger.error(f"Error processing version: {e}") + continue + + def dependencies(self) -> Generator[Dict[str, Any], None, None]: + """Transform PyPI dependency data into CHAI's dependency format.""" + dependency_type = self.db.select_dependency_type_by_name("runtime") + + for package_data in self._read_batch_files(): + try: + info = package_data.get("info") + try: + assert info, "Package info is missing" + except AssertionError as e: + self.logger.warn(str(e)) + continue + + name = info.get("name") + if not name: + continue + + pypi_id = f"pypi/{name}" + + # Process dependencies + requires_dist = info.get("requires_dist", []) + if requires_dist is None: + continue + + version = info.get("version") + if not version: + continue + + for req in requires_dist: + try: + # Parse dependency name and version + dep = self._parse_dependency(req) + + dep_name = dep[0] + dep_version = dep[1] + dep_version_range = dep[2] + + if not dep_name: + continue + + # Sometimes we don't get a version, so we need to get it from the database + if not dep_version: + # And getting the latest version + dep_version_from_db = self.db.select_latest_version_by_import_id(f"pypi/{dep_name}") + if dep_version_from_db: + dep_version = dep_version_from_db.version + else: + continue + + # Use the original version range for semver_range if available + # Otherwise use the clean version with an equals operator + semver_range = dep_version_range if dep_version_range else f"=={dep_version}" + + yield { + "version_id": f"pypi/{dep_name}-{dep_version}", + "import_id": pypi_id, + "semver_range": semver_range, + "dependency_type_id": dependency_type.id, + } + except Exception as e: + self.logger.error(f"Error processing dependency {req}: {e}") + continue + except Exception as e: + self.logger.error(f"Error processing package {package_data.get('info', {}).get('name')}: {e}") + continue + + def users(self) -> Generator[Dict[str, Any], None, None]: + """Skip user data as we can't get GitHub info from PyPI API.""" + if False: # This ensures it's a generator but never yields anything + yield {} + + def user_packages(self) -> Generator[Dict[str, Any], None, None]: + """Skip user-package relationships as we can't get GitHub info from PyPI API.""" + if False: + yield {} + + def user_versions(self) -> Generator[Dict[str, Any], None, None]: + """Skip user-version relationships as we can't get GitHub info from PyPI API.""" + if False: + yield {} + + def _parse_dependency(self, req: str) -> tuple[str, str, str]: + """Parse a dependency string into name, clean version, and version range. + + Examples: + "charset-normalizer (>=2,<4)" -> ("charset-normalizer", "2", ">=2,<4") + "idna (>=2.5,<4)" -> ("idna", "2.5", ">=2.5,<4") + "urllib3 (>=1.21.1,<3)" -> ("urllib3", "1.21.1", ">=1.21.1,<3") + "PySocks (>=1.5.6,!=1.5.7)" -> ("PySocks", "1.5.6", ">=1.5.6,!=1.5.7") + "pycparser ==2.21" -> ("pycparser", "2.21", "==2.21") + "pytest >=5.4.1,<6.0.0" -> ("pytest", "5.4.1", ">=5.4.1,<6.0.0") + "python-jose[cryptography] >=3.1.0,<4.0.0" -> ("python-jose", "3.1.0", ">=3.1.0,<4.0.0") + "GPUtil~=1.4.0" -> ("GPUtil", "1.4.0", "~=1.4.0") + "requests ~=2.25" -> ("requests", "2.25", "~=2.25") + "ipykernel<7.0.0,>=6.29.3" -> ("ipykernel", "6.29.3", "<7.0.0,>=6.29.3") + "pydantic>=2.10" -> ("pydantic", "2.10", ">=2.10") + """ + try: + # Remove any extra conditions after semicolon + if "; " in req: + req = req.split("; ")[0].strip() + + # First, handle extras in square brackets + name = req + if "[" in req and "]" in req: + name = req.split("[")[0].strip() + after_bracket = req[req.find("]")+1:].strip() + if after_bracket: + req = name + after_bracket + + # Handle parentheses format: "name (version)" + if "(" in req and ")" in req: + name = req.split(" (")[0].strip() + version_constraints = req.split("(")[1].rstrip(")").strip() + # Get both clean version and version range + clean_version = self._extract_version(version_constraints) + return name, clean_version, version_constraints.replace(" ", "") + + # Handle various version constraint formats + version_operators = [">=", "<=", "==", "!=", "~=", ">", "<", "="] + + # Find the first version operator in the string + operator_pos = len(req) + found_operator = None + + for op in version_operators: + pos = req.find(op) + if pos != -1 and pos < operator_pos: + operator_pos = pos + found_operator = op + + if found_operator: + # Split at the operator position + name = req[:operator_pos].strip() + version_constraints = req[operator_pos:].strip() + # Get both clean version and version range + clean_version = self._extract_version(version_constraints) + return name, clean_version, version_constraints.replace(" ", "") + + # If no version info found + return req.strip(), "", "" + + except Exception as e: + self.logger.error(f"Error parsing dependency string '{req}': {e}") + return req.strip(), "", "" + + def _extract_version(self, version_constraints: str) -> str: + """Extract the first valid version number from version constraints. + + Examples: + ">=2,<4" -> "2" + ">=2.5,<4" -> "2.5" + ">=1.21.1,<3" -> "1.21.1" + ">=1.5.6,!=1.5.7" -> "1.5.6" + "==2.21" -> "2.21" + ">=5.4.1,<6.0.0" -> "5.4.1" + ">=3.1.0,<4.0.0" -> "3.1.0" + "~=1.4.0" -> "1.4.0" + """ + # Remove all spaces + version_constraints = "".join(version_constraints.split()) + + # List of version operators to strip + operators = [">=", "<=", "==", "!=", "~=", ">", "<", "="] + + # Split by comma to handle multiple constraints + parts = version_constraints.split(",") + + for part in parts: + # Strip any operators + clean_version = part + for op in operators: + if clean_version.startswith(op): + clean_version = clean_version[len(op):] + + # If we have a valid version number (contains a digit), return it + if any(c.isdigit() for c in clean_version): + return clean_version + + return "" diff --git a/tests/unit/test_crates_transformer.py b/tests/unit/test_crates_transformer.py index 7844f465..bcf6d6f1 100644 --- a/tests/unit/test_crates_transformer.py +++ b/tests/unit/test_crates_transformer.py @@ -65,7 +65,7 @@ def test_versions_transform(self, transformer, mock_csv_reader): - Processing of optional fields """ test_data = { - "crate_id": "123", + "import_id": "123", "num": "1.0.0", "id": "456", "crate_size": "1000", @@ -81,7 +81,7 @@ def test_versions_transform(self, transformer, mock_csv_reader): assert len(versions) == 1 version = versions[0] - assert version["crate_id"] == "123" + assert version["import_id"] == "123" assert version["version"] == "1.0.0" assert version["import_id"] == "456" assert version["size"] == 1000 @@ -101,7 +101,7 @@ def test_dependencies_transform(self, transformer, mock_csv_reader): """ test_data = { "version_id": "456", - "crate_id": "789", + "import_id": "789", "req": "^1.0", "kind": "0", # normal dependency } @@ -113,7 +113,7 @@ def test_dependencies_transform(self, transformer, mock_csv_reader): dependency = dependencies[0] assert dependency["version_id"] == "456" - assert dependency["crate_id"] == "789" + assert dependency["import_id"] == "789" assert dependency["semver_range"] == "^1.0" assert dependency["dependency_type"] == DependencyType(0)