diff --git a/aibom/analyzer.py b/aibom/analyzer.py index 3335132..0b6092a 100644 --- a/aibom/analyzer.py +++ b/aibom/analyzer.py @@ -45,7 +45,13 @@ def visit_Call(self, node: ast.Call) -> Any: name = self._name_of(node.func) leaf = name.split(".")[-1] if leaf in MODEL_CLASS_HINTS: - self.models.append({"type": leaf, "model": self._arg_or_kw(node, "model", "model_name"), "source_file": str(self.file_path)}) + self.models.append( + { + "type": leaf, + "model": self._arg_or_kw(node, "model", "model_name"), + "source_file": str(self.file_path), + } + ) if leaf in TOOL_HINTS or "agent" in leaf.lower(): self.tools.append({"name": leaf, "source_file": str(self.file_path)}) if any(part in VECTORSTORE_HINTS for part in name.split(".")): @@ -68,9 +74,17 @@ def _name_of(self, node: ast.AST) -> str: def _arg_or_kw(self, node: ast.Call, *keys: str, default: str = "unknown") -> str: for kw in node.keywords: - if kw.arg in keys and isinstance(kw.value, ast.Constant) and isinstance(kw.value.value, str): + if ( + kw.arg in keys + and isinstance(kw.value, ast.Constant) + and isinstance(kw.value.value, str) + ): return kw.value.value - if node.args and isinstance(node.args[0], ast.Constant) and isinstance(node.args[0].value, str): + if ( + node.args + and isinstance(node.args[0], ast.Constant) + and isinstance(node.args[0].value, str) + ): return node.args[0].value return default diff --git a/aibom/bundle.py b/aibom/bundle.py index e94f0c9..1c420f2 100644 --- a/aibom/bundle.py +++ b/aibom/bundle.py @@ -1,6 +1,5 @@ from __future__ import annotations -import json from pathlib import Path from zipfile import ZIP_DEFLATED, ZipFile @@ -14,7 +13,9 @@ def build_manifest(files: dict[str, bytes]) -> dict[str, str]: return {name: sha256_bytes(content) for name, content in sorted(files.items())} -def create_bundle(aibom_path: Path, out_zip: Path, baseline_path: Path | None = None, compliance_md: str = "") -> Path: +def create_bundle( + aibom_path: Path, out_zip: Path, baseline_path: Path | None = None, compliance_md: str = "" +) -> Path: aibom = load_json(aibom_path) files: dict[str, bytes] = {} files["AIBOM.json"] = stable_json(aibom).encode("utf-8") diff --git a/aibom/cli.py b/aibom/cli.py index f6130a9..e4a1b32 100644 --- a/aibom/cli.py +++ b/aibom/cli.py @@ -13,7 +13,6 @@ from aibom.storage import load_json, persist_run from aibom.validation import AIBOMValidationException, validate_aibom - COMPLIANCE_STARTER = """# Starter Compliance Mapping\n\nThis mapping is a starter reference only and not legal advice.\n""" @@ -39,12 +38,15 @@ def cmd_generate(args: argparse.Namespace) -> int: _write_json(spdx_out, export_spdx(aibom)) if args.bundle_out: baseline = target / ".aibom" / "baseline.json" - create_bundle(out, Path(args.bundle_out).resolve(), baseline if baseline.exists() else None, COMPLIANCE_STARTER) + create_bundle( + out, + Path(args.bundle_out).resolve(), + baseline if baseline.exists() else None, + COMPLIANCE_STARTER, + ) return 0 - - def cmd_validate(args: argparse.Namespace) -> int: src = load_json(Path(args.input)) try: @@ -55,6 +57,7 @@ def cmd_validate(args: argparse.Namespace) -> int: print("OK: AIBOM validates against schema") return 0 + def cmd_export(args: argparse.Namespace) -> int: src = load_json(Path(args.input)) if args.format == "spdx-json": diff --git a/aibom/diffing.py b/aibom/diffing.py index 9d3b673..a7b747b 100644 --- a/aibom/diffing.py +++ b/aibom/diffing.py @@ -16,7 +16,11 @@ def diff_aibom(old: dict[str, Any], new: dict[str, Any]) -> dict[str, Any]: added = sorted([new_i[x] for x in new_i.keys() - old_i.keys()], key=lambda x: str(x)) removed = sorted([old_i[x] for x in old_i.keys() - new_i.keys()], key=lambda x: str(x)) changed = sorted( - [{"before": old_i[x], "after": new_i[x]} for x in new_i.keys() & old_i.keys() if old_i[x] != new_i[x]], + [ + {"before": old_i[x], "after": new_i[x]} + for x in new_i.keys() & old_i.keys() + if old_i[x] != new_i[x] + ], key=lambda x: str(x), ) out["added"][sec] = added diff --git a/aibom/exporters.py b/aibom/exporters.py index 616ff2b..b10d69e 100644 --- a/aibom/exporters.py +++ b/aibom/exporters.py @@ -6,11 +6,29 @@ def export_spdx(aibom: dict[str, Any]) -> dict[str, Any]: packages: list[dict[str, Any]] = [] for model in aibom.get("models", []): - packages.append({"SPDXID": f"SPDXRef-Model-{model['type']}", "name": model["type"], "versionInfo": model.get("model", "unknown")}) + packages.append( + { + "SPDXID": f"SPDXRef-Model-{model['type']}", + "name": model["type"], + "versionInfo": model.get("model", "unknown"), + } + ) for tool in aibom.get("tools", []): - packages.append({"SPDXID": f"SPDXRef-Tool-{tool['name']}", "name": tool["name"], "versionInfo": "unknown"}) + packages.append( + { + "SPDXID": f"SPDXRef-Tool-{tool['name']}", + "name": tool["name"], + "versionInfo": "unknown", + } + ) for ds in aibom.get("datasets", []): - packages.append({"SPDXID": f"SPDXRef-Dataset-{ds['type'].replace('.', '-')}", "name": ds["type"], "versionInfo": "unknown"}) + packages.append( + { + "SPDXID": f"SPDXRef-Dataset-{ds['type'].replace('.', '-')}", + "name": ds["type"], + "versionInfo": "unknown", + } + ) return { "spdxVersion": "SPDX-2.3", "dataLicense": "CC0-1.0", @@ -24,7 +42,18 @@ def export_spdx(aibom: dict[str, Any]) -> dict[str, Any]: def export_cyclonedx(aibom: dict[str, Any]) -> dict[str, Any]: components: list[dict[str, str]] = [] for model in aibom.get("models", []): - components.append({"type": "machine-learning-model", "name": model["type"], "version": model.get("model", "unknown")}) + components.append( + { + "type": "machine-learning-model", + "name": model["type"], + "version": model.get("model", "unknown"), + } + ) for fw in aibom.get("frameworks", []): components.append({"type": "library", "name": fw["name"], "version": "unknown"}) - return {"bomFormat": "CycloneDX", "specVersion": "1.5", "version": 1, "components": sorted(components, key=lambda x: x["name"])} + return { + "bomFormat": "CycloneDX", + "specVersion": "1.5", + "version": 1, + "components": sorted(components, key=lambda x: x["name"]), + } diff --git a/aibom_generator.py b/aibom_generator.py index 8797a79..ecf9f87 100644 --- a/aibom_generator.py +++ b/aibom_generator.py @@ -3,6 +3,5 @@ from aibom.cli import main - if __name__ == "__main__": raise SystemExit(main()) diff --git a/tests/test_cli.py b/tests/test_cli.py index bfb7980..3c20f4e 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -50,7 +50,9 @@ def test_golden_fixture_validates_against_schema() -> None: def test_validation_fixtures_cover_valid_and_invalid_cases() -> None: fixtures_dir = Path(__file__).parent / "fixtures" valid_doc = json.loads((fixtures_dir / "valid_aibom.json").read_text(encoding="utf-8")) - invalid_doc = json.loads((fixtures_dir / "invalid_aibom_missing_field.json").read_text(encoding="utf-8")) + invalid_doc = json.loads( + (fixtures_dir / "invalid_aibom_missing_field.json").read_text(encoding="utf-8") + ) validate_aibom(valid_doc) with pytest.raises(AIBOMValidationException) as exc: @@ -87,8 +89,9 @@ def test_cli_validate_command_success_and_failure(tmp_path: Path) -> None: assert "/" in bad.stderr - -def test_generate_fails_closed_before_writing_output(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: +def test_generate_fails_closed_before_writing_output( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: from aibom import cli as cli_module invalid_doc = generate_aibom(_fixture_project()) @@ -118,12 +121,18 @@ def test_export_spdx_deterministic() -> None: doc = generate_aibom(_fixture_project()) spdx = export_spdx(doc) assert spdx["spdxVersion"] == "SPDX-2.3" - assert [p["SPDXID"] for p in spdx["packages"]] == sorted([p["SPDXID"] for p in spdx["packages"]]) + assert [p["SPDXID"] for p in spdx["packages"]] == sorted( + [p["SPDXID"] for p in spdx["packages"]] + ) def test_diff_detects_additions() -> None: old = {"models": [], "tools": [], "datasets": []} - new = {"models": [{"type": "ChatOpenAI"}], "tools": [{"name": "initialize_agent"}], "datasets": []} + new = { + "models": [{"type": "ChatOpenAI"}], + "tools": [{"name": "initialize_agent"}], + "datasets": [], + } d = diff_aibom(old, new) assert len(d["added"]["models"]) == 1 assert len(d["added"]["tools"]) == 1 @@ -139,5 +148,7 @@ def test_bundle_contains_manifest(tmp_path: Path) -> None: def test_cli_version() -> None: - proc = subprocess.run([sys.executable, "-m", "aibom.cli", "--version"], capture_output=True, text=True, check=True) + proc = subprocess.run( + [sys.executable, "-m", "aibom.cli", "--version"], capture_output=True, text=True, check=True + ) assert "aibom" in proc.stdout