apisec-code-bolt 0.1.0__py3-none-any.whl

apisec_code_bolt/analysis/binding_tracker.py

@@ -0,0 +1,341 @@
+"""
+PyCG-style variable-to-type binding tracker with fixed-point propagation.
+
+Tracks what types each variable could be bound to so that method calls on
+dynamically-typed receivers can be resolved to concrete targets.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from .call_graph_types import TypeBinding
+
+
+class BindingTracker:
+    """
+    Tracks variable-to-type bindings for dynamic dispatch resolution.
+
+    This implements the core of the PyCG algorithm: we track what types
+    each variable could be bound to, so we can resolve method calls.
+
+    KEY IMPROVEMENT: Uses fixed-point iteration to propagate bindings
+    transitively (x = y means x has all types that y has).
+    """
+
+    def __init__(self):
+        # Per-function bindings: (func_qname, var_name) -> TypeBinding
+        self._local_bindings: dict[tuple[str, str], TypeBinding] = {}
+
+        # Module-level bindings: (file_path, var_name) -> TypeBinding
+        self._module_bindings: dict[tuple[Path, str], TypeBinding] = {}
+
+        # Class attribute bindings: (class_qname, attr_name) -> TypeBinding
+        self._class_bindings: dict[tuple[str, str], TypeBinding] = {}
+
+        # Return value bindings: func_qname -> possible return types/callables
+        self._return_bindings: dict[str, TypeBinding] = {}
+
+        # Protocol/ABC implementations: protocol_name -> [implementing_classes]
+        self._protocol_implementations: dict[str, set[str]] = {}
+
+        # Propagation tracking
+        self._propagation_complete = False
+
+    def add_assignment(
+        self,
+        variable: str,
+        assigned_type: str | None,
+        file_path: Path,
+        line: int,
+        scope_function: str | None = None,
+        scope_class: str | None = None,
+        source_variable: str | None = None,
+        is_callable: bool = False,
+        callable_target: str | None = None,
+    ) -> None:
+        """Record an assignment that binds a variable to a type."""
+        self._propagation_complete = False
+        location = (file_path, line)
+
+        if scope_function:
+            key = (scope_function, variable)
+            if key not in self._local_bindings:
+                self._local_bindings[key] = TypeBinding(variable=variable)
+            binding = self._local_bindings[key]
+        elif scope_class:
+            key = (scope_class, variable)
+            if key not in self._class_bindings:
+                self._class_bindings[key] = TypeBinding(variable=variable)
+            binding = self._class_bindings[key]
+        else:
+            key = (file_path, variable)
+            if key not in self._module_bindings:
+                self._module_bindings[key] = TypeBinding(variable=variable)
+            binding = self._module_bindings[key]
+
+        if assigned_type:
+            binding.possible_types.add(assigned_type)
+        if source_variable:
+            binding.source_variables.add(source_variable)
+        if is_callable:
+            binding.holds_callable = True
+            if callable_target:
+                binding.callable_targets.add(callable_target)
+
+        binding.assignment_locations.append(location)
+
+    def add_self_binding(
+        self,
+        function_qname: str,
+        class_qname: str,
+        param_name: str = "self",
+    ) -> None:
+        """Bind self/cls parameter to its enclosing class."""
+        self._propagation_complete = False
+        key = (function_qname, param_name)
+        binding = TypeBinding(
+            variable=param_name,
+            is_self=param_name == "self",
+            is_cls=param_name == "cls",
+            enclosing_class=class_qname,
+        )
+        binding.possible_types.add(class_qname)
+        self._local_bindings[key] = binding
+
+    def add_parameter(
+        self,
+        param_name: str,
+        function_qname: str,
+        type_annotation: str | None = None,
+    ) -> None:
+        """Record a function parameter (could be any compatible type)."""
+        self._propagation_complete = False
+        key = (function_qname, param_name)
+        binding = TypeBinding(
+            variable=param_name,
+            is_parameter=True,
+            parameter_type_annotation=type_annotation,
+        )
+
+        # If we have a type annotation, that's the possible type
+        if type_annotation:
+            binding.possible_types.add(type_annotation)
+            # Check if it's a protocol/ABC
+            if type_annotation in self._protocol_implementations:
+                binding.possible_types.update(self._protocol_implementations[type_annotation])
+
+        self._local_bindings[key] = binding
+
+    def add_import(
+        self,
+        local_name: str,
+        import_source: str,
+        file_path: Path,
+        line: int,
+        is_star: bool = False,
+        star_exports: list[str] | None = None,
+    ) -> None:
+        """Record an import binding."""
+        self._propagation_complete = False
+        key = (file_path, local_name)
+        binding = TypeBinding(
+            variable=local_name,
+            is_import=True,
+            import_source=import_source,
+        )
+        binding.assignment_locations.append((file_path, line))
+        self._module_bindings[key] = binding
+
+        # Handle star imports
+        if is_star and star_exports:
+            for name in star_exports:
+                star_key = (file_path, name)
+                star_binding = TypeBinding(
+                    variable=name,
+                    is_import=True,
+                    import_source=f"{import_source}.{name}",
+                )
+                self._module_bindings[star_key] = star_binding
+
+    def add_return_binding(
+        self,
+        function_qname: str,
+        return_type: str | None = None,
+        returns_callable: bool = False,
+        callable_target: str | None = None,
+    ) -> None:
+        """Track what a function returns."""
+        self._propagation_complete = False
+        if function_qname not in self._return_bindings:
+            self._return_bindings[function_qname] = TypeBinding(
+                variable=f"<return:{function_qname}>"
+            )
+        binding = self._return_bindings[function_qname]
+
+        if return_type:
+            binding.possible_types.add(return_type)
+        if returns_callable:
+            binding.holds_callable = True
+            if callable_target:
+                binding.callable_targets.add(callable_target)
+
+    def add_protocol_implementation(
+        self,
+        protocol: str,
+        implementer: str,
+    ) -> None:
+        """Record that a class implements a protocol/ABC."""
+        self._propagation_complete = False
+        if protocol not in self._protocol_implementations:
+            self._protocol_implementations[protocol] = set()
+        self._protocol_implementations[protocol].add(implementer)
+
+    def propagate(self, max_iterations: int = 100) -> int:
+        """
+        Perform fixed-point iteration to propagate bindings transitively.
+
+        If x = y, then x should have all types that y has.
+
+        Returns: Number of iterations performed
+        """
+        if self._propagation_complete:
+            return 0
+
+        iterations = 0
+        changed = True
+
+        while changed and iterations < max_iterations:
+            changed = False
+            iterations += 1
+
+            # Propagate in local bindings
+            for key, binding in list(self._local_bindings.items()):
+                func_qname, var_name = key
+                for source_var in list(binding.source_variables):
+                    source_key = (func_qname, source_var)
+                    if source_key in self._local_bindings:
+                        source_binding = self._local_bindings[source_key]
+                        old_size = len(binding.possible_types)
+                        binding.possible_types.update(source_binding.possible_types)
+                        binding.callable_targets.update(source_binding.callable_targets)
+                        if source_binding.holds_callable:
+                            binding.holds_callable = True
+                        if len(binding.possible_types) > old_size:
+                            changed = True
+
+            # Propagate in module bindings
+            for key, binding in list(self._module_bindings.items()):
+                file_path, var_name = key
+                for source_var in list(binding.source_variables):
+                    source_key = (file_path, source_var)
+                    if source_key in self._module_bindings:
+                        source_binding = self._module_bindings[source_key]
+                        old_size = len(binding.possible_types)
+                        binding.possible_types.update(source_binding.possible_types)
+                        binding.callable_targets.update(source_binding.callable_targets)
+                        if source_binding.holds_callable:
+                            binding.holds_callable = True
+                        if len(binding.possible_types) > old_size:
+                            changed = True
+
+            # Propagate return values for call assignments
+            # If x = func(), x should have the return types of func
+            for _key, binding in list(self._local_bindings.items()):
+                for callable_name in list(binding.source_variables):
+                    if callable_name in self._return_bindings:
+                        return_binding = self._return_bindings[callable_name]
+                        old_size = len(binding.possible_types)
+                        binding.possible_types.update(return_binding.possible_types)
+                        binding.callable_targets.update(return_binding.callable_targets)
+                        if return_binding.holds_callable:
+                            binding.holds_callable = True
+                        if len(binding.possible_types) > old_size:
+                            changed = True
+
+        self._propagation_complete = True
+        return iterations
+
+    def get_possible_types(
+        self,
+        variable: str,
+        file_path: Path,
+        scope_function: str | None = None,
+    ) -> set[str]:
+        """Get possible types for a variable in a given scope."""
+        if not self._propagation_complete:
+            self.propagate()
+
+        types: set[str] = set()
+
+        # Check local scope first
+        if scope_function:
+            key = (scope_function, variable)
+            if key in self._local_bindings:
+                types.update(self._local_bindings[key].possible_types)
+
+        # Check module scope
+        key = (file_path, variable)
+        if key in self._module_bindings:
+            binding = self._module_bindings[key]
+            types.update(binding.possible_types)
+            if binding.is_import and binding.import_source:
+                types.add(binding.import_source)
+
+        return types
+
+    def get_callable_targets(
+        self,
+        variable: str,
+        file_path: Path,
+        scope_function: str | None = None,
+    ) -> set[str]:
+        """Get callable targets for a variable that holds a function."""
+        if not self._propagation_complete:
+            self.propagate()
+
+        targets: set[str] = set()
+
+        if scope_function:
+            key = (scope_function, variable)
+            if key in self._local_bindings:
+                binding = self._local_bindings[key]
+                targets.update(binding.callable_targets)
+
+        key = (file_path, variable)
+        if key in self._module_bindings:
+            binding = self._module_bindings[key]
+            targets.update(binding.callable_targets)
+
+        return targets
+
+    def get_binding(
+        self,
+        variable: str,
+        file_path: Path,
+        scope_function: str | None = None,
+    ) -> TypeBinding | None:
+        """Get the binding for a variable."""
+        if scope_function:
+            key = (scope_function, variable)
+            if key in self._local_bindings:
+                return self._local_bindings[key]
+
+        key = (file_path, variable)
+        return self._module_bindings.get(key)
+
+    def get_return_types(self, function_qname: str) -> set[str]:
+        """Get the return types of a function."""
+        if function_qname in self._return_bindings:
+            return self._return_bindings[function_qname].possible_types
+        return set()
+
+    def get_returned_callables(self, function_qname: str) -> set[str]:
+        """Get callables returned by a function."""
+        if function_qname in self._return_bindings:
+            return self._return_bindings[function_qname].callable_targets
+        return set()
+
+    def get_protocol_implementers(self, protocol: str) -> set[str]:
+        """Get all classes that implement a protocol/ABC."""
+        return self._protocol_implementations.get(protocol, set())