"""Conservative lexical data-flow features for Python snippets. The analyzer resolves simple names within lexical scopes. Attribute storage, subscript keys, aliases beyond direct ``a = b`` assignments, imported behavior, dynamic dispatch, and runtime-dependent control flow remain unresolved. """ import re from collections import Counter from .call_graph import analyze_local_call_graph from .node_types import ( ASSIGNMENT, ATTRIBUTE, AUGMENTED_ASSIGNMENT, CLASS_DEFINITION, FUNCTION_DEFINITION, IDENTIFIER, IMPORT_FROM_STATEMENT, IMPORT_STATEMENT, ) from .parser import get_node_text, walk SCOPE_TYPES = {FUNCTION_DEFINITION, CLASS_DEFINITION, "lambda"} TARGET_CONTAINERS = { "list_pattern", "pattern_list", "tuple_pattern", } def _simple_target_names(node) -> list[str]: if node is None: return [] if node.type == IDENTIFIER: return [get_node_text(node)] if node.type in TARGET_CONTAINERS: return [ get_node_text(item) for item in walk(node) if item.type == IDENTIFIER ] if node.type in {"list_splat_pattern", "dictionary_splat_pattern"}: return [ get_node_text(item) for item in walk(node) if item.type == IDENTIFIER ] return [] def _parameter_names(function) -> list[str]: parameters = function.child_by_field_name("parameters") if parameters is None: return [] names = [] for child in parameters.named_children: if child.type == IDENTIFIER: names.append(get_node_text(child)) continue name = child.child_by_field_name("name") if name is not None and name.type == IDENTIFIER: names.append(get_node_text(name)) continue if child.type in {"list_splat_pattern", "dictionary_splat_pattern"}: names.extend(_simple_target_names(child)) return names def _imported_names(node) -> list[str]: text = get_node_text(node) if node.type == IMPORT_STATEMENT: body = re.sub(r"^\s*import\s+", "", text) names = [] for item in body.split(","): parts = item.strip().split() if not parts: continue names.append(parts[-1] if len(parts) >= 3 and parts[-2] == "as" else parts[0].split(".")[0]) return names match = re.match(r"\s*from\s+[.\w]+\s+import\s+(.+)", text) if not match: return [] names = [] for item in match.group(1).strip("()").split(","): parts = item.strip().split() if not parts or parts[0] == "*": continue names.append(parts[-1] if len(parts) >= 3 and parts[-2] == "as" else parts[0]) return names def _scope_local_names(scope) -> set[str]: names = set(_parameter_names(scope)) if scope.type == FUNCTION_DEFINITION else set() def visit(node): if node is not scope and node.type in SCOPE_TYPES: name = node.child_by_field_name("name") if name is not None: names.add(get_node_text(name)) return if node.type in {ASSIGNMENT, AUGMENTED_ASSIGNMENT}: names.update(_simple_target_names(node.child_by_field_name("left"))) elif node.type in {"for_statement", "named_expression"}: names.update(_simple_target_names(node.child_by_field_name("left"))) elif node.type == "with_item": alias = node.child_by_field_name("alias") if alias is not None: names.update(_simple_target_names(alias)) elif node.type == "except_clause": value = node.child_by_field_name("value") if value is not None and value.type == "as_pattern": names.update( get_node_text(item) for item in walk(value) if item.type == IDENTIFIER ) elif node.type in {IMPORT_STATEMENT, IMPORT_FROM_STATEMENT}: names.update(_imported_names(node)) for child in node.named_children: visit(child) visit(scope) return names def _assignment_names(node) -> set[str]: names = set() def visit(current): if current is not node and current.type in SCOPE_TYPES: return if current.type in {ASSIGNMENT, AUGMENTED_ASSIGNMENT}: names.update( _simple_target_names(current.child_by_field_name("left")) ) elif current.type == "for_statement": names.update( _simple_target_names(current.child_by_field_name("left")) ) for child in current.named_children: visit(child) visit(node) return names def _branch_merge_count(root_node) -> int: merges = 0 for node in walk(root_node): arms = [] if node.type == "if_statement": consequence = node.child_by_field_name("consequence") if consequence is not None: arms.append(_assignment_names(consequence)) for child in node.named_children: if child.type not in {"elif_clause", "else_clause"}: continue consequence = child.child_by_field_name("consequence") body = consequence or next( ( item for item in child.named_children if item.type == "block" ), None, ) if body is not None: arms.append(_assignment_names(body)) elif node.type == "match_statement": arms = [ _assignment_names(case) for case in walk(node) if case.type == "case_clause" ] if len(arms) >= 2: counts = Counter(name for arm in arms for name in arm) merges += sum(count >= 2 for count in counts.values()) return merges class _ScopeAnalyzer: def __init__(self, scope, ancestor_names: set[str]): self.scope = scope self.ancestor_names = ancestor_names self.local_names = _scope_local_names(scope) self.definitions = {} self.next_definition = 0 self.depths = {} self.local_edges = 0 self.cross_scope = 0 self.unresolved_reads = 0 self.alias_assignments = 0 self.max_depth = 0 self.nested_scopes = [] def define(self, name: str, dependency_depths=()): definition = self.next_definition self.next_definition += 1 depth = 1 + max(dependency_depths, default=0) self.definitions[name] = definition self.depths[definition] = depth self.max_depth = max(self.max_depth, depth) def read(self, name: str) -> int | None: definition = self.definitions.get(name) if definition is not None: self.local_edges += 1 return self.depths[definition] if name not in self.local_names and name in self.ancestor_names: self.cross_scope += 1 return None self.unresolved_reads += 1 return None def consume_expression(self, node) -> list[int]: if node is None: return [] if node.type in SCOPE_TYPES: return [] if node.type == IDENTIFIER: depth = self.read(get_node_text(node)) return [depth] if depth is not None else [] if node.type == ATTRIBUTE: return self.consume_expression(node.child_by_field_name("object")) if node.type == "call": depths = [] function = node.child_by_field_name("function") if function is not None and function.type == ATTRIBUTE: depths.extend( self.consume_expression( function.child_by_field_name("object") ) ) arguments = node.child_by_field_name("arguments") if arguments is not None: for child in arguments.named_children: if child.type == "keyword_argument": value = child.child_by_field_name("value") depths.extend(self.consume_expression(value)) else: depths.extend(self.consume_expression(child)) return depths if node.type == "keyword_argument": return self.consume_expression(node.child_by_field_name("value")) depths = [] for child in node.named_children: depths.extend(self.consume_expression(child)) return depths def visit(self, node): if node is not self.scope and node.type in SCOPE_TYPES: name = node.child_by_field_name("name") if name is not None: self.define(get_node_text(name)) self.nested_scopes.append(node) return if node.type in {IMPORT_STATEMENT, IMPORT_FROM_STATEMENT}: for name in _imported_names(node): self.define(name) return if node.type == ASSIGNMENT: left = node.child_by_field_name("left") right = node.child_by_field_name("right") dependencies = self.consume_expression(right) targets = _simple_target_names(left) if ( len(targets) == 1 and right is not None and right.type == IDENTIFIER ): self.alias_assignments += 1 for name in targets: self.define(name, dependencies) if not targets: self.consume_expression(left) return if node.type == AUGMENTED_ASSIGNMENT: left = node.child_by_field_name("left") dependencies = self.consume_expression(left) dependencies.extend( self.consume_expression(node.child_by_field_name("right")) ) for name in _simple_target_names(left): self.define(name, dependencies) return if node.type == "for_statement": dependencies = self.consume_expression( node.child_by_field_name("right") ) for name in _simple_target_names( node.child_by_field_name("left") ): self.define(name, dependencies) body = node.child_by_field_name("body") if body is not None: self.visit(body) alternative = node.child_by_field_name("alternative") if alternative is not None: self.visit(alternative) return if node.type == "with_item": dependencies = self.consume_expression( node.child_by_field_name("value") ) alias = node.child_by_field_name("alias") for name in _simple_target_names(alias): self.define(name, dependencies) return if node.type == "call": self.consume_expression(node) return if node.type == IDENTIFIER: self.consume_expression(node) return for child in node.named_children: self.visit(child) def analyze(self): for name in _parameter_names(self.scope): self.define(name) body = ( self.scope.child_by_field_name("body") if self.scope.type in SCOPE_TYPES else self.scope ) if body is not None: for child in body.named_children: self.visit(child) return self def _mutation_counts(root_node) -> tuple[int, int]: attributes = 0 subscripts = 0 for node in walk(root_node): if node.type not in { ASSIGNMENT, AUGMENTED_ASSIGNMENT, "delete_statement", }: continue target = node.child_by_field_name("left") if target is None and node.type == "delete_statement": targets = node.named_children else: targets = [target] if target is not None else [] for item in targets: attributes += sum( descendant.type == ATTRIBUTE for descendant in walk(item) ) subscripts += sum( descendant.type == "subscript" for descendant in walk(item) ) return attributes, subscripts def analyze_data_flow(root_node, unresolved_call_count: int = 0) -> dict: """Extract lexical dependency and uncertainty metrics. Def-use edges only connect reads to earlier simple-name definitions in the same scope. Cross-scope reads are counted but not connected into dependency chains because runtime closure/global state is not modeled. """ pending = [(root_node, set())] local_edges = 0 cross_scope = 0 unresolved_reads = 0 aliases = 0 max_depth = 0 while pending: scope, ancestor_names = pending.pop() analyzer = _ScopeAnalyzer(scope, ancestor_names).analyze() local_edges += analyzer.local_edges cross_scope += analyzer.cross_scope unresolved_reads += analyzer.unresolved_reads aliases += analyzer.alias_assignments max_depth = max(max_depth, analyzer.max_depth) nested_ancestors = ancestor_names | analyzer.local_names pending.extend( (nested, nested_ancestors) for nested in analyzer.nested_scopes ) attribute_mutations, subscript_mutations = _mutation_counts(root_node) total_observations = ( local_edges + cross_scope + unresolved_reads + unresolved_call_count ) unresolved = unresolved_reads + unresolved_call_count return { "local_def_use_edge_count": local_edges, "max_dataflow_chain_depth": max_depth, "cross_scope_flow_count": cross_scope, "branch_merge_count": _branch_merge_count(root_node), "alias_assignment_count": aliases, "attribute_mutation_count": attribute_mutations, "subscript_mutation_count": subscript_mutations, "unresolved_call_count": unresolved_call_count, "unresolved_flow_ratio": round( unresolved / max(total_observations, 1), 3, ), } def extract_data_flow(root_node) -> dict: """Return the standalone data-flow signal group for a parsed snippet.""" call_graph = analyze_local_call_graph(root_node) return analyze_data_flow( root_node, unresolved_call_count=call_graph["unresolved_local_call_count"], )