test_normalize_python.py

"""Python normalizer: universal IR extraction from tree-sitter ASTs."""

from __future__ import annotations

import textwrap

from sciogen.ir import EdgeKind, FileIR, Node, NodeKind
from sciogen.normalize import get_normalizer


def normalize(source: str, path: str = "mod.py") -> FileIR:
    src = textwrap.dedent(source).lstrip("\n").encode()
    ir = FileIR(path=path, language="python", sha256="x", size=len(src), mtime=0.0)
    ir.nodes.append(Node(id=path, kind=NodeKind.FILE, name=path, file=path, language="python"))
    get_normalizer("python").normalize(ir, src)
    return ir


def node_by_id(ir: FileIR, node_id: str) -> Node:
    return next(n for n in ir.nodes if n.id == node_id)


def test_class_and_methods_extracted():
    ir = normalize('''
        class AuthService:
            """Authenticates users."""

            def login(self, username: str, password: str) -> bool:
                """Log a user in."""
                return True
    ''')
    cls = node_by_id(ir, "mod.py:AuthService")
    assert cls.kind == NodeKind.CLASS
    assert cls.docstring == "Authenticates users."
    method = node_by_id(ir, "mod.py:AuthService.login")
    assert method.kind == NodeKind.FUNCTION
    assert method.params == ["self", "username: str", "password: str"]
    assert method.return_type == "bool"
    defines = {(e.src, e.dst) for e in ir.edges if e.kind == EdgeKind.DEFINES}
    assert ("mod.py", "mod.py:AuthService") in defines
    assert ("mod.py:AuthService", "mod.py:AuthService.login") in defines


def test_imports_all_forms():
    ir = normalize('''
        import hashlib
        import os.path as osp
        from src.models.user import UserModel
        from .hashing import hash_password as hp
    ''')
    by_alias = {i.alias: i for i in ir.imports}
    assert by_alias["hashlib"].is_module
    assert by_alias["osp"].target == "os.path"
    assert by_alias["UserModel"].module == "src.models.user"
    assert by_alias["hp"].target == ".hashing.hash_password"


def test_calls_recorded_with_context():
    ir = normalize('''
        class Service:
            def run(self):
                self.helper()
                validate(self.data)
                registry.get_client().send()

            def helper(self):
                pass
    ''')
    calls = {c.callee for c in ir.calls if c.caller_id == "mod.py:Service.run"}
    assert calls == {"self.helper", "validate", "registry.get_client", "?.send"}
    assert all(c.enclosing_class == "Service" for c in ir.calls)


def test_mutations_self_and_global():
    ir = normalize('''
        CACHE = {}

        class Store:
            def save(self, value):
                self.items = [value]

        def reset():
            global CACHE
            CACHE = {}
    ''')
    targets = {(m.function_id, m.target) for m in ir.mutations}
    assert ("mod.py:Store.save", "self.items") in targets
    assert ("mod.py:reset", "CACHE") in targets


def test_complexity_counts_branches():
    ir = normalize('''
        def branchy(x):
            if x > 0:
                for i in range(x):
                    if i % 2 and i % 3:
                        x += 1
            return x
    ''')
    assert node_by_id(ir, "mod.py:branchy").complexity == 5  # if + for + if + and + 1


def test_typed_dict_becomes_type_node_with_fields():
    ir = normalize('''
        from typing import TypedDict

        class UserPayload(TypedDict):
            name: str
            email: str
    ''')
    t = node_by_id(ir, "mod.py:UserPayload")
    assert t.kind == NodeKind.TYPE
    assert t.fields == ["name", "email"]


def test_long_function_gets_chunks():
    body = "\n".join(f"    step_{i} = compute_{i}()" for i in range(60))
    ir = normalize(f"def pipeline():\n{body}\n")
    chunks = [n for n in ir.nodes if n.kind == NodeKind.CHUNK]
    assert len(chunks) >= 2
    assert all(c.parent == "mod.py:pipeline" for c in chunks)
    assert all(e.kind == EdgeKind.PART_OF for e in ir.edges if e.src.startswith("mod.py:pipeline#"))
    # chunk embedding text is a normalized summary, not raw code
    assert "step_1" in chunks[0].embedding_text
    assert "=" not in chunks[0].embedding_text


def test_nested_definitions_and_module_guards():
    ir = normalize('''
        try:
            import fast_json as json
        except ImportError:
            import json

        if True:
            def guarded():
                pass

        def outer():
            def inner():
                pass
            return inner
    ''')
    ids = {n.id for n in ir.nodes}
    assert {"mod.py:guarded", "mod.py:outer", "mod.py:outer.inner"} <= ids
    assert {i.alias for i in ir.imports} == {"json"}