diff --git a/wtismycode-core/src/package_classifier.rs b/wtismycode-core/src/package_classifier.rs index ae4883c..016653c 100644 --- a/wtismycode-core/src/package_classifier.rs +++ b/wtismycode-core/src/package_classifier.rs @@ -356,7 +356,7 @@ lazy_static::lazy_static! { (PackageCategory::Database, vec![ "sqlalchemy", "psycopg2", "psycopg", "asyncpg", "pymongo", "mongoengine", "peewee", "tortoise", "databases", - "alembic", "pymysql", "opensearch", "elasticsearch", + "alembic", "pymysql", "opensearch", "opensearchpy", "elasticsearch", "motor", "beanie", "odmantic", "sqlmodel", "piccolo", "edgedb", "cassandra", "clickhouse_driver", "sqlite3", "neo4j", "arango", "influxdb", "timescaledb", @@ -396,7 +396,19 @@ lazy_static::lazy_static! { ]), (PackageCategory::Logging, vec![ "loguru", "structlog", "sentry_sdk", "watchtower", - "python_json_logger", "colorlog", "rich", + "python_json_logger", "colorlog", "rich", "prometheus_client", + ]), + (PackageCategory::AiMl, vec![ + "pyannote", "soundfile", "librosa", "audioread", "webrtcvad", + ]), + (PackageCategory::Queue, vec![ + "aiormq", + ]), + (PackageCategory::Http, vec![ + "pydantic_settings", "pydantic_extra_types", "fastapi_mail", + ]), + (PackageCategory::Database, vec![ + "peewee_async", "peewee_migrate", ]), ]; } diff --git a/wtismycode-core/src/python_analyzer.rs b/wtismycode-core/src/python_analyzer.rs index 2bef3ca..09c1236 100644 --- a/wtismycode-core/src/python_analyzer.rs +++ b/wtismycode-core/src/python_analyzer.rs @@ -686,8 +686,21 @@ impl PythonAnalyzer { self.compute_metrics(&mut project_model)?; // Classify all imports using PackageClassifier + // Collect all known project module names to filter from integrations + let project_modules: std::collections::HashSet = modules.iter() + .map(|m| { + let mod_path = self.compute_module_path(&m.path); + mod_path.split('.').next().unwrap_or(&mod_path).to_lowercase() + }) + .collect(); + let all_imports: Vec = modules.iter() .flat_map(|m| m.imports.iter().map(|i| i.module_name.clone())) + .filter(|import| { + let top = import.split('.').next().unwrap_or(import).to_lowercase(); + // Skip imports that are project's own modules + !project_modules.contains(&top) + }) .collect(); let cache_dir = if self.config.caching.enabled {