Add Python requirement extraction from docstrings and comments

graphify/extract.py

@@ -40,7 +40,10 @@
 from graphify.paths import disambiguate_ambiguous_candidates
 
 _RECURSION_LIMIT = 10_000
-
+_REQUIREMENT_ID_RE = re.compile(
+    r"\b[A-Z]{2,}-\d+\b",
+    re.IGNORECASE,
+)
 # Language built-in globals that AST may classify as call targets when used as
 # constructors or coercion functions (e.g. String(x), Number(x), Boolean(x)).
 # Without this filter they become god-nodes accumulating spurious edges from
@@ -96,7 +99,7 @@ def _csharp_namespace_id(dotted_name: str) -> str:
 
 SEMANTIC_RELATIONS = frozenset({
     "inherits", "implements", "mixes_in", "embeds", "references",
-    "calls", "imports", "imports_from", "re_exports", "contains", "method",
+    "calls", "imports", "imports_from", "re_exports", "contains", "method","satisfies",
 })
 
 REFERENCE_CONTEXTS = frozenset({
@@ -5737,7 +5740,241 @@ def _is_autogenerated_python(source: bytes) -> bool:
         return True
     return False
 
+# Requirement IDs like SAFE-04, DRIVE-05, AUTH-10
+_REQUIREMENT_ID_RE = re.compile(
+    r"\b([A-Z][A-Z0-9_]*-\d+)\b"
+)
+
+
+def _extract_requirement_ids(text: str) -> list[str]:
+    """
+    Extract requirement identifiers from text.
+
+    Examples:
+        SAFE-04
+        AUTH-12
+        PERF-001
+
+    Returns:
+        List of requirement IDs.
+    """
+    if not text:
+        return []
+
+    return [
+        match.group(1)
+        for match in _REQUIREMENT_ID_RE.finditer(text)
+    ]
+def _extract_python_requirements(path: Path, result: dict) -> None:
+    """
+    Post-pass: extract requirement IDs from Python docstrings/comments.
+
+    Creates:
+        code_symbol ---- satisfies ----> requirement
+
+    Mutates result in-place.
+    """
+    try:
+        import tree_sitter_python as tspython
+        from tree_sitter import Language, Parser
+
+        language = Language(tspython.language())
+        parser = Parser(language)
+
+        source = path.read_bytes()
+        tree = parser.parse(source)
+        root = tree.root_node
+
+    except Exception:
+        return
+
+    stem = _file_stem(path)
+    str_path = str(path)
+
+    nodes = result["nodes"]
+    edges = result["edges"]
+
+    seen_ids = {n["id"] for n in nodes}
+    seen_edges = {
+        (e["source"], e["target"], e["relation"])
+        for e in edges
+    }
+
+    file_nid = _make_id(str(path))
+
+
+    def _get_docstring(body_node):
+        if not body_node:
+            return None
+
+        for child in body_node.children:
+            if child.type == "expression_statement":
+                for sub in child.children:
+                    if sub.type in ("string", "concatenated_string"):
+                        text = source[
+                            sub.start_byte:sub.end_byte
+                        ].decode("utf-8", errors="replace")
+
+                        return (
+                            text
+                            .strip("\"'")
+                            .strip('"""')
+                            .strip("'''")
+                        )
+
+            break
+
+        return None
+
+
+    def _add_requirement(req_id, line, symbol_id):
+
+        req_node_id = _make_id("requirement", req_id)
+
+        if req_node_id not in seen_ids:
+            seen_ids.add(req_node_id)
+
+            nodes.append({
+                "id": req_node_id,
+                "label": req_id,
+                "file_type": "concept",
+                "source_file": str_path,
+                "source_location": f"L{line}",
+            })
+
+
+        edge_key = (
+            symbol_id,
+            req_node_id,
+            "satisfies"
+        )
+
+        if edge_key not in seen_edges:
+            seen_edges.add(edge_key)
 
+            edges.append({
+                "source": symbol_id,
+                "target": req_node_id,
+                "relation": "satisfies",
+                "confidence": "EXTRACTED",
+                "source_file": str_path,
+                "source_location": f"L{line}",
+                "weight": 1.0,
+            })
+
+
+    def _scan_text(text, line, symbol_id):
+
+        if not text:
+            return
+
+        req_ids = _extract_requirement_ids(text)
+
+        for req_id in req_ids:
+            _add_requirement(
+                req_id,
+                line,
+                symbol_id
+            )
+
+
+    def walk(node, parent_nid):
+
+        if node.type == "class_definition":
+
+            name_node = node.child_by_field_name("name")
+
+            body = node.child_by_field_name("body")
+
+            if name_node and body:
+
+                class_name = source[
+                    name_node.start_byte:name_node.end_byte
+                ].decode()
+
+                class_nid = _make_id(
+                    stem,
+                    class_name
+                )
+
+                doc = _get_docstring(body)
+
+                if doc:
+                    _scan_text(
+                        doc,
+                        body.start_point[0] + 1,
+                        class_nid
+                    )
+
+                for child in body.children:
+                    walk(child, class_nid)
+
+            return
+
+
+        if node.type == "function_definition":
+
+            name_node = node.child_by_field_name("name")
+
+            body = node.child_by_field_name("body")
+
+            if name_node and body:
+
+                func_name = source[
+                    name_node.start_byte:name_node.end_byte
+                ].decode()
+
+
+                func_nid = (
+                    _make_id(parent_nid, func_name)
+                    if parent_nid != file_nid
+                    else _make_id(stem, func_name)
+                )
+
+
+                doc = _get_docstring(body)
+
+                if doc:
+                    _scan_text(
+                        doc,
+                        body.start_point[0] + 1,
+                        func_nid
+                    )
+
+            return
+
+
+        for child in node.children:
+            walk(child, parent_nid)
+
+
+    # Scan docstrings
+    walk(root, file_nid)
+
+
+    # Scan comments
+    source_text = source.decode(
+        "utf-8",
+        errors="replace"
+    )
+
+    for lineno, line in enumerate(
+        source_text.splitlines(),
+        start=1
+    ):
+
+        if "#" in line:
+
+            comment = line.split("#",1)[1]
+
+            req_ids = _extract_requirement_ids(comment)
+
+            for req_id in req_ids:
+                _add_requirement(
+                    req_id,
+                    lineno,
+                    file_nid
+                )
 def _extract_python_rationale(path: Path, result: dict) -> None:
     """Post-pass: extract docstrings and rationale comments from Python source.
     Mutates result in-place by appending to result['nodes'] and result['edges'].
@@ -5849,6 +6086,7 @@ def extract_python(path: Path) -> dict:
     result = _extract_generic(path, _PYTHON_CONFIG)
     if "error" not in result:
         _extract_python_rationale(path, result)
+        _extract_python_requirements(path, result)
     return result
 
 

tests/test_rationale.py

@@ -33,7 +33,33 @@ def process():
     rationale = [n for n in result["nodes"] if n.get("file_type") == "rationale"]
     assert any("chunked" in n["label"] for n in rationale)
 
+def test_requirement_docstring_creates_satisfies_edge(tmp_path):
+    path = _write_py(tmp_path, '''
+        def process():
+            """
+            SAFE-04: Validate all inputs before processing.
+            """
+            pass
+    ''')
+
+    result = extract_python(path)
+
+    req_nodes = [
+        n for n in result["nodes"]
+        if n.get("file_type") == "concept"
+    ]
+
+    assert len(req_nodes) == 1
+    assert req_nodes[0]["label"] == "SAFE-04"
+
+    satisfies_edges = [
+        e for e in result["edges"]
+        if e.get("relation") == "satisfies"
+    ]
 
+    assert len(satisfies_edges) == 1
+    assert satisfies_edges[0]["confidence"] == "EXTRACTED"
+    assert satisfies_edges[0]["weight"] == 1.0
 def test_class_docstring_extracted(tmp_path):
     path = _write_py(tmp_path, '''
         class Cache: