diff --git a/tests/test_generate_popular_node_signatures.py b/tests/test_generate_popular_node_signatures.py index 9387c08..2096c0f 100644 --- a/tests/test_generate_popular_node_signatures.py +++ b/tests/test_generate_popular_node_signatures.py @@ -12,6 +12,10 @@ from tools.generate_popular_node_signatures import ( class StaticExtractionTests(unittest.TestCase): + def _normalise_generated_at(self, text): + parsed = json.loads(text) + return text.replace(parsed["generated_at"], "") + def test_normalise_input_spec_reduces_combo_lists(self): self.assertEqual("COMBO", normalise_input_spec((["nearest", "bilinear"],))) self.assertEqual("IMAGE", normalise_input_spec(("IMAGE",))) @@ -99,15 +103,105 @@ NODE_CLASS_MAPPINGS = { self.assertEqual({}, result["nodes"]) self.assertEqual("no_static_nodes", result["pack"]["status"]) + def test_skips_unparseable_python_files_and_extracts_static_nodes(self): + good_source = ''' +class GoodNode: + RETURN_TYPES = ("IMAGE",) + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + }, + } + + +NODE_CLASS_MAPPINGS = { + "GoodNode": GoodNode, +} +''' + with tempfile.TemporaryDirectory() as tmp: + Path(tmp, "bad.py").write_bytes(b"class Bad:\xff\n") + Path(tmp, "good.py").write_text(textwrap.dedent(good_source), encoding="utf-8") + result = extract_repo_signatures( + Path(tmp), + { + "id": "mixed-pack", + "title": "Mixed Pack", + "repository": "https://github.com/example/mixed-pack", + "rank": 1, + }, + ) + + self.assertIn("GoodNode", result["nodes"]) + self.assertEqual("ok", result["pack"]["status"]) + + def test_unsupported_reassignment_invalidates_static_env_value(self): + source = ''' +def build_inputs(): + return {"required": {"image": ("IMAGE",)}} + + +INPUTS = { + "required": { + "image": ("IMAGE",), + }, +} +INPUTS = build_inputs() + + +class StaleEnvNode: + RETURN_TYPES = ("IMAGE",) + + @classmethod + def INPUT_TYPES(cls): + return INPUTS + + +NODE_CLASS_MAPPINGS = { + "StaleEnvNode": StaleEnvNode, +} +''' + with tempfile.TemporaryDirectory() as tmp: + Path(tmp, "__init__.py").write_text(textwrap.dedent(source), encoding="utf-8") + result = extract_repo_signatures( + Path(tmp), + { + "id": "stale-env-pack", + "title": "Stale Env Pack", + "repository": "https://github.com/example/stale-env-pack", + "rank": 1, + }, + ) + + self.assertEqual({}, result["nodes"]) + self.assertEqual("no_static_nodes", result["pack"]["status"]) + def test_write_artifact_is_deterministic(self): with tempfile.TemporaryDirectory() as tmp: - out = Path(tmp, "popular_node_signatures.json") + out_one = Path(tmp, "one.json") + out_two = Path(tmp, "two.json") write_artifact( - out, - sources={"manager_url": "https://example.invalid/manager.json", "limit": 1}, + out_one, + sources={ + "manager_url": "https://example.invalid/manager.json", + "limit": 1, + "registry": {"z": "last", "a": "first"}, + }, packs={ - "b-pack": {"id": "b-pack", "title": "B Pack", "status": "ok"}, - "a-pack": {"id": "a-pack", "title": "A Pack", "status": "ok"}, + "b-pack": { + "id": "b-pack", + "title": "B Pack", + "status": "ok", + "metadata": {"z": 2, "a": 1}, + }, + "a-pack": { + "id": "a-pack", + "title": "A Pack", + "status": "ok", + "metadata": {"z": 4, "a": 3}, + }, }, nodes={ "BNode": { @@ -115,7 +209,7 @@ NODE_CLASS_MAPPINGS = { "display": "B Node", "pack": "b-pack", "repository": "https://github.com/example/b-pack", - "inputs": {}, + "inputs": {"zeta": "FLOAT", "alpha": "IMAGE"}, "required": [], "outputs": ["IMAGE"], "output_names": ["image"], @@ -126,7 +220,7 @@ NODE_CLASS_MAPPINGS = { "display": "A Node", "pack": "a-pack", "repository": "https://github.com/example/a-pack", - "inputs": {}, + "inputs": {"zeta": "FLOAT", "alpha": "IMAGE"}, "required": [], "outputs": ["IMAGE"], "output_names": ["image"], @@ -134,10 +228,59 @@ NODE_CLASS_MAPPINGS = { }, }, ) - parsed = json.loads(out.read_text(encoding="utf-8")) + write_artifact( + out_two, + sources={ + "registry": {"a": "first", "z": "last"}, + "limit": 1, + "manager_url": "https://example.invalid/manager.json", + }, + packs={ + "a-pack": { + "metadata": {"a": 3, "z": 4}, + "status": "ok", + "title": "A Pack", + "id": "a-pack", + }, + "b-pack": { + "metadata": {"a": 1, "z": 2}, + "status": "ok", + "title": "B Pack", + "id": "b-pack", + }, + }, + nodes={ + "ANode": { + "confidence": "static_exact", + "output_names": ["image"], + "outputs": ["IMAGE"], + "required": [], + "inputs": {"alpha": "IMAGE", "zeta": "FLOAT"}, + "repository": "https://github.com/example/a-pack", + "pack": "a-pack", + "display": "A Node", + "type": "ANode", + }, + "BNode": { + "confidence": "static_exact", + "output_names": ["image"], + "outputs": ["IMAGE"], + "required": [], + "inputs": {"alpha": "IMAGE", "zeta": "FLOAT"}, + "repository": "https://github.com/example/b-pack", + "pack": "b-pack", + "display": "B Node", + "type": "BNode", + }, + }, + ) + text_one = out_one.read_text(encoding="utf-8") + text_two = out_two.read_text(encoding="utf-8") + parsed = json.loads(text_one) self.assertEqual(["a-pack", "b-pack"], list(parsed["packs"])) self.assertEqual(["ANode", "BNode"], list(parsed["nodes"])) + self.assertEqual(self._normalise_generated_at(text_one), self._normalise_generated_at(text_two)) if __name__ == "__main__": diff --git a/tools/generate_popular_node_signatures.py b/tools/generate_popular_node_signatures.py index 136a9c2..f90c56f 100644 --- a/tools/generate_popular_node_signatures.py +++ b/tools/generate_popular_node_signatures.py @@ -42,9 +42,11 @@ def _collect_module_env(tree): continue if len(stmt.targets) != 1 or not isinstance(stmt.targets[0], ast.Name): continue + name = stmt.targets[0].id try: - env[stmt.targets[0].id] = _literal(stmt.value, env) + env[name] = _literal(stmt.value, env) except UnsupportedStaticExpression: + env.pop(name, None) continue return env @@ -181,14 +183,23 @@ def _python_files(repo_dir): yield Path(root, filename) +def _parse_python_file(path): + try: + return ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + except UnicodeDecodeError: + try: + return ast.parse(path.read_text(encoding="utf-8", errors="ignore"), filename=str(path)) + except SyntaxError: + return None + except SyntaxError: + return None + + def extract_repo_signatures(repo_dir, pack_meta): nodes = {} for path in sorted(_python_files(repo_dir)): - try: - tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) - except UnicodeDecodeError: - tree = ast.parse(path.read_text(encoding="utf-8", errors="ignore"), filename=str(path)) - except SyntaxError: + tree = _parse_python_file(path) + if tree is None: continue env = _collect_module_env(tree) mappings = _node_class_mappings(tree, env) @@ -213,13 +224,21 @@ def extract_repo_signatures(repo_dir, pack_meta): return {"pack": pack, "nodes": nodes} +def _sorted_json_value(value): + if isinstance(value, dict): + return {key: _sorted_json_value(value[key]) for key in sorted(value)} + if isinstance(value, list): + return [_sorted_json_value(item) for item in value] + return value + + def write_artifact(path, sources, packs, nodes): payload = { "schema_version": SCHEMA_VERSION, "generated_at": datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z"), - "sources": sources, - "packs": {key: packs[key] for key in sorted(packs)}, - "nodes": {key: nodes[key] for key in sorted(nodes)}, + "sources": _sorted_json_value(sources), + "packs": _sorted_json_value(packs), + "nodes": _sorted_json_value(nodes), } path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(payload, indent=2, sort_keys=False) + "\n", encoding="utf-8")