pycheck-neuron 1.0.0__tar.gz → 1.0.1__tar.gz

{pycheck_neuron-1.0.0/pycheck_neuron.egg-info → pycheck_neuron-1.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pycheck-neuron
-Version: 1.0.0
+Version: 1.0.1
 Summary: 32-rule static AST analyzer catching temporal lookahead bias, causal confusion, and uncertainty bugs in Python ML code. Zero dependencies.
 Author: NEURON Team
 Author-email: Fayo Ibrahim <fayo@neuron-lab.org>

{pycheck_neuron-1.0.0 → pycheck_neuron-1.0.1}/pycheck/__init__.py

@@ -3,7 +3,7 @@ PyCheck — NEURON ML Safety Analyzer
 Catches temporal leaks, causal confusion, and uncertainty bugs in Python ML scripts.
 """
 
-__version__ = "0.2.0"
+__version__ = "1.0.1"
 
 from pycheck.analyzer import analyze_file, format_output
 from pycheck.rules import ALL_RULES, RULES_BY_CODE

{pycheck_neuron-1.0.0 → pycheck_neuron-1.0.1}/pycheck/analyzer.py

@@ -109,14 +109,18 @@ def analyze_file(filepath: str) -> Tuple[List[dict], List[str]]:
     source_lines = source.split('\n')
     tree = ast.parse(source, filename=filepath)
     
-    # Phase 1: Rule-based AST analysis
-    analyzer = PyCheckAnalyzer(source_lines, filepath)
-    analyzer.visit(tree)
-    
-    # Phase 2: Data flow taint analysis
+    # Phase 1: Data flow taint & provenance analysis
     tracker = TaintTracker()
     flow_diagnostics = tracker.analyze(tree, source_lines)
     
+    # Phase 2: Rule-based AST analysis (with provenance populated in ctx)
+    analyzer = PyCheckAnalyzer(source_lines, filepath)
+    analyzer.ctx.tainted_vars = {k: v.reason for k, v in tracker.tainted_vars.items()}
+    analyzer.ctx.train_vars = tracker.train_vars
+    analyzer.ctx.test_vars = tracker.test_vars
+    analyzer.ctx.test_contaminated_vars = tracker.test_contaminated_vars
+    analyzer.visit(tree)
+    
     # Combine and deduplicate
     all_diagnostics = analyzer.diagnostics + flow_diagnostics
     

pycheck_neuron-1.0.1/pycheck/flow.py

@@ -0,0 +1,547 @@
+"""
+PyCheck Data Flow Engine — Taint Propagation Analysis
+Tracks how data flows through assignments, function returns, and method chains
+to detect when future-tainted data reaches training/prediction sinks.
+"""
+
+import ast
+from typing import Dict, Set, List, Optional, Tuple
+from pycheck.rules import Diagnostic
+
+
+class TaintSource:
+    """Represents a source of tainted (future-leaked) data."""
+    def __init__(self, var_name: str, reason: str, line: int, col: int):
+        self.var_name = var_name
+        self.reason = reason
+        self.line = line
+        self.col = col
+
+
+class TaintTracker:
+    """
+    Tracks taint propagation and partition data provenance through Python AST.
+    
+    Taint sources: shift(-n), bfill(), future column access, etc.
+    Provenance sources: train_test_split, temporal slicing, cross-validation splits.
+    Contamination: concatenating or reassigning train data with test partitions.
+    Sinks: .fit(), .fit_transform(), model training calls.
+    
+    Emits:
+      - F001: Future-tainted data flows into sinks
+      - F002: Test partition or contaminated data flows into training sinks
+    """
+    
+    def __init__(self):
+        self.tainted_vars: Dict[str, TaintSource] = {}
+        self.diagnostics: List[Diagnostic] = []
+        self.constants: Dict[str, any] = {}
+        
+        # Partition provenance tracking
+        self.train_vars: Set[str] = set()
+        self.test_vars: Set[str] = set()
+        self.test_contaminated_vars: Dict[str, str] = {}
+        self.aliases: Dict[str, str] = {}
+        self.functions: Dict[str, ast.FunctionDef] = {}
+        self.train_markers = ('_train', 'train_', 'x_train', 'y_train', 'x_tr', 'y_tr', '_trn', 'trn_', 'train_df', 'train_data', 'train_set')
+        self.test_markers = ('_test', 'test_', 'x_test', 'y_test', 'x_val', 'y_val', '_val', 'val_', '_eval', 'eval_', 'holdout', 'validation', 'test_df', 'test_data', 'test_set')
+    
+    def analyze(self, tree: ast.AST, source_lines: List[str]) -> List[Diagnostic]:
+        """Run taint and provenance analysis on a full AST."""
+        self.tainted_vars.clear()
+        self.diagnostics.clear()
+        self.constants.clear()
+        self.train_vars.clear()
+        self.test_vars.clear()
+        self.test_contaminated_vars.clear()
+        self.aliases.clear()
+        self.functions.clear()
+        
+        # Two-pass analysis:
+        # Pass 1: Identify taint & provenance sources, follow reassignments and combinations
+        # Pass 2: Detect sensitive sinks
+        self._collect_taints(tree)
+        self._check_sinks(tree)
+        
+        return self.diagnostics
+    
+    def _iter_statements(self, node: ast.AST):
+        """Yield statements recursively in source code execution order."""
+        if hasattr(node, 'body') and isinstance(node.body, list):
+            for stmt in node.body:
+                yield stmt
+                yield from self._iter_statements(stmt)
+        if hasattr(node, 'orelse') and isinstance(node.orelse, list):
+            for stmt in node.orelse:
+                yield stmt
+                yield from self._iter_statements(stmt)
+        if hasattr(node, 'handlers') and isinstance(node.handlers, list):
+            for h in node.handlers:
+                yield from self._iter_statements(h)
+        if hasattr(node, 'finalbody') and isinstance(node.finalbody, list):
+            for stmt in node.finalbody:
+                yield stmt
+                yield from self._iter_statements(stmt)
+
+    def _extract_constant(self, expr):
+        if isinstance(expr, ast.Constant) and isinstance(expr.value, (int, float)):
+            return expr.value
+        if isinstance(expr, ast.UnaryOp) and isinstance(expr.op, ast.USub):
+            if isinstance(expr.operand, ast.Constant) and isinstance(expr.operand.value, (int, float)):
+                return -expr.operand.value
+        return None
+
+    def _is_negative_period(self, expr) -> bool:
+        if isinstance(expr, ast.UnaryOp) and isinstance(expr.op, ast.USub):
+            return True
+        if isinstance(expr, ast.Constant) and isinstance(expr.value, (int, float)) and expr.value < 0:
+            return True
+        if isinstance(expr, ast.Name):
+            val = self.constants.get(expr.id)
+            if val is not None and isinstance(val, (int, float)) and val < 0:
+                return True
+        return False
+
+    def _collect_taints(self, tree: ast.AST):
+        """Pass 1: Walk the AST, index helper functions, track taint and provenance."""
+        # Index helper functions defined in the AST
+        for node in ast.walk(tree):
+            if isinstance(node, ast.FunctionDef):
+                self.functions[node.name] = node
+
+        # Extract literal constants
+        for node in ast.walk(tree):
+            if isinstance(node, ast.Assign):
+                for target in node.targets:
+                    if isinstance(target, ast.Name):
+                        val = self._extract_constant(node.value)
+                        if val is not None:
+                            self.constants[target.id] = val
+
+        # Linear traversal of statements for data-flow & provenance propagation
+        for stmt in self._iter_statements(tree):
+            if isinstance(stmt, ast.Assign):
+                # 1. Check for split calls (train_test_split)
+                if isinstance(stmt.value, ast.Call):
+                    call_name = self._get_call_name(stmt.value)
+                    if call_name == 'train_test_split':
+                        for target in stmt.targets:
+                            if isinstance(target, (ast.Tuple, ast.List)):
+                                names = [self._get_var_name(e) for e in target.elts]
+                                if len(names) >= 4:
+                                    if names[0]: self.train_vars.add(names[0])
+                                    if names[1]: self.test_vars.add(names[1])
+                                    if names[2]: self.train_vars.add(names[2])
+                                    if names[3]: self.test_vars.add(names[3])
+                                elif len(names) >= 2:
+                                    if names[0]: self.train_vars.add(names[0])
+                                    if names[1]: self.test_vars.add(names[1])
+
+                # 2. Check for temporal slicing splits: df[:split_point] vs df[split_point:]
+                if isinstance(stmt.value, ast.Subscript):
+                    sl = stmt.value.slice
+                    if isinstance(sl, ast.Slice):
+                        for target in stmt.targets:
+                            t_name = self._get_var_name(target)
+                            if t_name:
+                                if sl.lower is None or (isinstance(sl.lower, ast.Constant) and sl.lower.value == 0):
+                                    if sl.upper is not None:
+                                        self.train_vars.add(t_name)
+                                elif sl.upper is None and sl.lower is not None:
+                                    self.test_vars.add(t_name)
+                    elif isinstance(sl, ast.Name):
+                        for target in stmt.targets:
+                            t_name = self._get_var_name(target)
+                            if t_name:
+                                if sl.id in self.train_vars or any(m in sl.id.lower() for m in self.train_markers):
+                                    self.train_vars.add(t_name)
+                                elif sl.id in self.test_vars or any(m in sl.id.lower() for m in self.test_markers):
+                                    self.test_vars.add(t_name)
+
+                # 3. Register naming evidence
+                for target in stmt.targets:
+                    t_name = self._get_var_name(target)
+                    if t_name:
+                        t_lower = t_name.lower()
+                        if any(m in t_lower for m in self.train_markers) and t_name not in self.test_contaminated_vars:
+                            self.train_vars.add(t_name)
+                        elif any(m in t_lower for m in self.test_markers):
+                            self.test_vars.add(t_name)
+
+                # 4. Standard future-shift / bfill taint check
+                self._check_assignment_taint(stmt)
+
+                # 5. Check test contamination propagation
+                contamination = self._expr_contains_test(stmt.value)
+                for target in stmt.targets:
+                    t_name = self._get_var_name(target)
+                    if not t_name:
+                        continue
+                    if contamination:
+                        self.test_contaminated_vars[t_name] = contamination
+                        self.train_vars.discard(t_name)
+                    else:
+                        # Follow aliases
+                        src_name = self._get_var_name(stmt.value)
+                        if src_name:
+                            self.aliases[t_name] = src_name
+                            if src_name in self.test_contaminated_vars:
+                                self.test_contaminated_vars[t_name] = f"alias of contaminated '{src_name}' ({self.test_contaminated_vars[src_name]})"
+                                self.train_vars.discard(t_name)
+                            elif src_name in self.test_vars:
+                                self.test_vars.add(t_name)
+                            elif src_name in self.train_vars and t_name not in self.test_contaminated_vars:
+                                self.train_vars.add(t_name)
+
+            elif isinstance(stmt, ast.AugAssign):
+                self._check_augassign_taint(stmt)
+                contamination = self._expr_contains_test(stmt.value)
+                t_name = self._get_var_name(stmt.target)
+                if t_name and contamination:
+                    self.test_contaminated_vars[t_name] = f"augmented assignment ({type(stmt.op).__name__}) with {contamination}"
+                    self.train_vars.discard(t_name)
+
+            elif isinstance(stmt, ast.For):
+                if isinstance(stmt.iter, ast.Call) and self._get_call_name(stmt.iter) in ('split',):
+                    if isinstance(stmt.target, (ast.Tuple, ast.List)) and len(stmt.target.elts) == 2:
+                        idx0 = self._get_var_name(stmt.target.elts[0])
+                        idx1 = self._get_var_name(stmt.target.elts[1])
+                        if idx0: self.train_vars.add(idx0)
+                        if idx1: self.test_vars.add(idx1)
+
+    def _check_assignment_taint(self, node: ast.Assign):
+        """Check if an assignment creates or propagates a taint."""
+        # Get the taint from the RHS
+        taint_reason = self._expr_taint(node.value)
+        if taint_reason is None:
+            return
+        
+        # Mark all LHS targets as tainted
+        for target in node.targets:
+            var_name = self._get_var_name(target)
+            if var_name:
+                self.tainted_vars[var_name] = TaintSource(
+                    var_name, taint_reason, node.lineno, node.col_offset)
+    
+    def _check_augassign_taint(self, node: ast.AugAssign):
+        """Check if an augmented assignment (+=, etc.) propagates taint."""
+        taint_reason = self._expr_taint(node.value)
+        if taint_reason:
+            var_name = self._get_var_name(node.target)
+            if var_name:
+                self.tainted_vars[var_name] = TaintSource(
+                    var_name, taint_reason, node.lineno, node.col_offset)
+    
+    def _expr_taint(self, node: ast.AST) -> Optional[str]:
+        """
+        Determine if an expression produces tainted data.
+        Returns the taint reason string, or None if clean.
+        """
+        if isinstance(node, ast.Call):
+            return self._call_taint(node)
+        
+        if isinstance(node, ast.Name):
+            # Propagation: if RHS is a tainted variable, the assignment is tainted
+            if node.id in self.tainted_vars:
+                return f"derived from tainted variable '{node.id}' ({self.tainted_vars[node.id].reason})"
+        
+        if isinstance(node, ast.BinOp):
+            # Binary operations propagate taint from either operand
+            left_taint = self._expr_taint(node.left)
+            right_taint = self._expr_taint(node.right)
+            return left_taint or right_taint
+        
+        if isinstance(node, ast.Subscript):
+            # Subscript propagates taint from the value being indexed
+            val_taint = self._expr_taint(node.value)
+            if val_taint:
+                return val_taint
+            # Check if indexing a tainted column name
+            if isinstance(node.slice, ast.Constant) and isinstance(node.slice.value, str):
+                col_lower = node.slice.value.lower()
+                # Only flag column names that unambiguously refer to future data.
+                # 'target' and 'label' are intentionally excluded — they are standard
+                # ML label column names and produce false positives when accessed on a
+                # correctly split train/test dataset.
+                future_cols = {'future_price', 'future_return', 'next_close', 'next_price',
+                               'next_day', 'tomorrow', 'price_change_next_day', 'future_avg_price'}
+                if any(f in col_lower for f in future_cols):
+                    return f"accessing future-named column '{node.slice.value}'"
+        
+        if isinstance(node, ast.Attribute):
+            return self._expr_taint(node.value)
+        
+        return None
+    
+    def _call_taint(self, node: ast.Call) -> Optional[str]:
+        """Check if a function call produces tainted output."""
+        if isinstance(node.func, ast.Attribute):
+            method = node.func.attr
+            
+            # Direct taint sources
+            if method == 'shift':
+                for arg in node.args:
+                    if self._is_negative_period(arg):
+                        return "shift() with negative period (future access)"
+                for kw in node.keywords:
+                    if kw.arg == 'periods' and self._is_negative_period(kw.value):
+                        return "shift() with negative period (future access)"
+            
+            if method in ('bfill', 'backfill'):
+                return f".{method}() backfills from future values"
+            
+            if method == 'fillna':
+                for kw in node.keywords:
+                    if kw.arg == 'method' and isinstance(kw.value, ast.Constant):
+                        if kw.value.value in ('bfill', 'backfill'):
+                            return f".fillna(method='{kw.value.value}') backfills from future"
+            
+            if method == 'diff':
+                for arg in node.args:
+                    if self._is_negative_period(arg):
+                        return ".diff() with negative period (future access)"
+                for kw in node.keywords:
+                    if kw.arg == 'periods' and self._is_negative_period(kw.value):
+                        return ".diff() with negative period (future access)"
+            
+            # Taint propagation through method chains
+            obj_taint = self._expr_taint(node.func.value)
+            if obj_taint:
+                # Most method calls on tainted data produce tainted output
+                non_taint_methods = {'shape', 'dtype', 'dtypes', 'columns', 'index',
+                                     'describe', 'info', 'head', 'tail', 'len', '__len__'}
+                if method not in non_taint_methods:
+                    return obj_taint
+        
+        if isinstance(node.func, ast.Name):
+            # Check if calling a function with tainted arguments
+            for arg in node.args:
+                arg_taint = self._expr_taint(arg)
+                if arg_taint:
+                    return f"function receives tainted input: {arg_taint}"
+        
+        # R2-9 fix: Detect getattr(obj, 'shift')(-1) and similar patterns
+        # Pattern: call(getattr(obj, 'shift'), -1) => obj.shift(-1)
+        if isinstance(node.func, ast.Call):
+            inner_call = node.func
+            if isinstance(inner_call.func, ast.Name) and inner_call.func.id == 'getattr':
+                if len(inner_call.args) >= 2 and isinstance(inner_call.args[1], ast.Constant):
+                    attr_name = inner_call.args[1].value
+                    if attr_name == 'shift':
+                        for arg in node.args:
+                            if self._is_negative_period(arg):
+                                return "getattr(..., 'shift')() with negative period (future access)"
+                        for kw in node.keywords:
+                            if kw.arg == 'periods' and self._is_negative_period(kw.value):
+                                return "getattr(..., 'shift')() with negative period (future access)"
+                    elif attr_name == 'diff':
+                        for arg in node.args:
+                            if self._is_negative_period(arg):
+                                return "getattr(..., 'diff')() with negative period (future access)"
+                        for kw in node.keywords:
+                            if kw.arg == 'periods' and self._is_negative_period(kw.value):
+                                return "getattr(..., 'diff')() with negative period (future access)"
+                    elif attr_name in ('bfill', 'backfill'):
+                        return f"getattr(..., '{attr_name}')() backfills from future values"
+        
+        return None
+    
+    def _check_sinks(self, tree: ast.AST):
+        """Pass 2: Check if tainted or test-contaminated data flows into sensitive sinks."""
+        for node in ast.walk(tree):
+            if not isinstance(node, ast.Call):
+                continue
+            
+            if isinstance(node.func, ast.Attribute):
+                method = node.func.attr
+                sinks = {'fit', 'fit_transform', 'fit_predict', 'train',
+                         'partial_fit', 'predict', 'score'}
+                
+                if method in sinks:
+                    # Check positional arguments for future taint (F001)
+                    for arg in node.args:
+                        arg_taint = self._expr_taint(arg)
+                        if arg_taint:
+                            self.diagnostics.append(Diagnostic(
+                                node.lineno, node.col_offset, "error", "F001",
+                                f"Tainted data flows into .{method}(): {arg_taint}",
+                                "Ensure training/prediction data is free of future information"))
+                    # Check keyword arguments for future taint (F001)
+                    for kw in node.keywords:
+                        kw_taint = self._expr_taint(kw.value)
+                        if kw_taint:
+                            kw_name = kw.arg if kw.arg else "**kwargs"
+                            self.diagnostics.append(Diagnostic(
+                                node.lineno, node.col_offset, "error", "F001",
+                                f"Tainted data flows into .{method}() via keyword '{kw_name}': {kw_taint}",
+                                "Ensure training/prediction data is free of future information"))
+
+                # Check F002: Data contamination & test partitions in training sinks
+                training_sinks = {'fit', 'fit_transform', 'fit_predict', 'train', 'partial_fit'}
+                if method in training_sinks:
+                    obj_name = self._get_obj_name(node.func.value)
+                    safe_objects = {'fig', 'figure', 'ax', 'axes', 'plt', 'layout',
+                                    'canvas', 'widget', 'app', 'db', 'session', 'socket',
+                                    'parser', 'formatter', 'logger', 'handler'}
+                    if obj_name and obj_name.lower() in safe_objects:
+                        continue
+                    
+                    for arg in node.args:
+                        var_name = self._get_var_name(arg)
+                        if var_name:
+                            base_var = var_name.split('.')[0]
+                            if base_var in self.test_contaminated_vars:
+                                reason = self.test_contaminated_vars[base_var]
+                                self.diagnostics.append(Diagnostic(
+                                    node.lineno, node.col_offset, "error", "F002",
+                                    f"Data contamination: '{base_var}' contains test partition data ({reason}) but is passed to .{method}()",
+                                    "Ensure data passed to .{method}() contains only genuine training data without test partitions"))
+                            elif base_var in self.test_vars or (base_var not in self.train_vars and any(m in base_var.lower() for m in self.test_markers)):
+                                self.diagnostics.append(Diagnostic(
+                                    node.lineno, node.col_offset, "error", "F002",
+                                    f"Test data passed to training sink: '{base_var}' is a test/validation partition but is passed to .{method}()",
+                                    "Call .{method}() only on training data, use .transform() or .predict() on test data"))
+                    
+                    for kw in node.keywords:
+                        var_name = self._get_var_name(kw.value)
+                        if var_name:
+                            base_var = var_name.split('.')[0]
+                            kw_name = kw.arg if kw.arg else "**kwargs"
+                            if base_var in self.test_contaminated_vars:
+                                reason = self.test_contaminated_vars[base_var]
+                                self.diagnostics.append(Diagnostic(
+                                    node.lineno, node.col_offset, "error", "F002",
+                                    f"Data contamination: keyword '{kw_name}' receives '{base_var}' which contains test partition data ({reason}) in .{method}()",
+                                    "Ensure data passed to .{method}() contains only genuine training data without test partitions"))
+                            elif base_var in self.test_vars or (base_var not in self.train_vars and any(m in base_var.lower() for m in self.test_markers)):
+                                self.diagnostics.append(Diagnostic(
+                                    node.lineno, node.col_offset, "error", "F002",
+                                    f"Test data passed to training sink: keyword '{kw_name}' receives test partition '{base_var}' in .{method}()",
+                                    "Call .{method}() only on training data, use .transform() or .predict() on test data"))
+
+    def _expr_contains_test(self, expr: Optional[ast.AST]) -> Optional[str]:
+        """Check if an expression incorporates test partition data or contaminated variables."""
+        if expr is None:
+            return None
+        
+        # 1. Direct variable name
+        if isinstance(expr, ast.Name):
+            var_id = expr.id
+            if var_id in self.test_contaminated_vars:
+                return f"contaminated variable '{var_id}' ({self.test_contaminated_vars[var_id]})"
+            if var_id in self.test_vars:
+                return f"test partition '{var_id}'"
+            if var_id not in self.train_vars and any(m in var_id.lower() for m in self.test_markers):
+                return f"test partition '{var_id}'"
+            return None
+        
+        # 2. Function / Method Calls
+        if isinstance(expr, ast.Call):
+            call_name = self._get_call_name(expr)
+            # Concatenation: pd.concat([X_train, X_test]), np.concatenate, np.vstack
+            if call_name in ('concat', 'concatenate', 'vstack', 'hstack'):
+                if expr.args:
+                    arg0 = expr.args[0]
+                    if isinstance(arg0, (ast.List, ast.Tuple)):
+                        for elt in arg0.elts:
+                            reason = self._expr_contains_test(elt)
+                            if reason:
+                                return f"concatenation ({call_name}) including {reason}"
+                    else:
+                        reason = self._expr_contains_test(arg0)
+                        if reason:
+                            return f"concatenation ({call_name}) including {reason}"
+            
+            # Method calls: .append(), .merge(), .join()
+            if isinstance(expr.func, ast.Attribute):
+                method = expr.func.attr
+                if method in ('append', 'merge', 'join'):
+                    obj_reason = self._expr_contains_test(expr.func.value)
+                    if obj_reason:
+                        return f".{method}() on {obj_reason}"
+                    for arg in expr.args:
+                        arg_reason = self._expr_contains_test(arg)
+                        if arg_reason:
+                            return f".{method}() with {arg_reason}"
+            
+            # Helper functions defined in this module
+            if isinstance(expr.func, ast.Name):
+                func_name = expr.func.id
+                if func_name in self.functions:
+                    f_def = self.functions[func_name]
+                    param_names = [a.arg for a in f_def.args.args]
+                    contaminated_params = []
+                    for idx, arg in enumerate(expr.args):
+                        reason = self._expr_contains_test(arg)
+                        if reason:
+                            p_name = param_names[idx] if idx < len(param_names) else f"arg{idx}"
+                            contaminated_params.append((p_name, reason))
+                    
+                    if contaminated_params:
+                        # Check return statements in f_def
+                        ret_nodes = [n for n in ast.walk(f_def) if isinstance(n, ast.Return) and n.value is not None]
+                        for p_name, reason in contaminated_params:
+                            for ret in ret_nodes:
+                                ret_names = {n.id for n in ast.walk(ret.value) if isinstance(n, ast.Name)}
+                                if p_name in ret_names:
+                                    return f"helper function '{func_name}()' returning {reason}"
+                        p_name, reason = contaminated_params[0]
+                        return f"helper function '{func_name}()' receiving {reason}"
+        
+        # 3. Binary operations: a + b
+        if isinstance(expr, ast.BinOp):
+            left_r = self._expr_contains_test(expr.left)
+            if left_r:
+                return left_r
+            return self._expr_contains_test(expr.right)
+        
+        # 4. List / Tuple: [a, b]
+        if isinstance(expr, (ast.List, ast.Tuple)):
+            for elt in expr.elts:
+                r = self._expr_contains_test(elt)
+                if r:
+                    return r
+        
+        # 5. Subscript: obj[...]
+        if isinstance(expr, ast.Subscript):
+            val_r = self._expr_contains_test(expr.value)
+            if val_r:
+                return val_r
+            slice_r = self._expr_contains_test(expr.slice)
+            if slice_r:
+                return f"indexed by {slice_r}"
+        
+        # 6. Attribute: obj.attr
+        if isinstance(expr, ast.Attribute):
+            return self._expr_contains_test(expr.value)
+        
+        return None
+
+    def _get_call_name(self, node: ast.AST) -> str:
+        if isinstance(node, ast.Call):
+            if isinstance(node.func, ast.Name):
+                return node.func.id
+            if isinstance(node.func, ast.Attribute):
+                return node.func.attr
+        return ""
+
+    def _get_obj_name(self, node: ast.AST) -> Optional[str]:
+        if isinstance(node, ast.Name):
+            return node.id
+        if isinstance(node, ast.Attribute):
+            return node.attr
+        return None
+
+    def _get_var_name(self, node: ast.AST) -> Optional[str]:
+        """Extract variable name from an assignment target or expression."""
+        if isinstance(node, ast.Name):
+            return node.id
+        if isinstance(node, ast.Subscript):
+            base = self._get_var_name(node.value)
+            if isinstance(node.slice, ast.Constant) and isinstance(node.slice.value, str):
+                return f"{base}.{node.slice.value}" if base else node.slice.value
+            return base
+        if isinstance(node, ast.Attribute):
+            base = self._get_var_name(node.value)
+            return f"{base}.{node.attr}" if base else node.attr
+        return None

{pycheck_neuron-1.0.0 → pycheck_neuron-1.0.1}/pycheck/rules.py

@@ -49,6 +49,9 @@ class AnalysisContext:
         self.loop_var: Optional[str] = None
         self.assignments: Dict[str, ast.AST] = {}  # var_name -> assigned node
         self.constants: Dict[str, any] = {}  # var_name -> literal constant value
+        self.train_vars: Set[str] = set()
+        self.test_vars: Set[str] = set()
+        self.test_contaminated_vars: Dict[str, str] = {}
 
 
 # ═══════════════════════════════════════════════════════════════
@@ -311,15 +314,25 @@ class T010_FitBeforeSplit(Rule):
     severity = "error"
     category = "TemporalLeak"
 
-    def _is_train_data(self, node):
+    def _is_train_data(self, node, ctx=None):
         """Recursively check if an AST node represents training data."""
+        var_name = self._get_obj_name(node)
+        if var_name and ctx is not None:
+            base_var = var_name.split('.')[0]
+            if hasattr(ctx, 'test_contaminated_vars') and base_var in ctx.test_contaminated_vars:
+                return False
+            if hasattr(ctx, 'test_vars') and base_var in ctx.test_vars:
+                return False
+            if hasattr(ctx, 'train_vars') and base_var in ctx.train_vars:
+                return True
+
         train_markers = ('_train', 'train_', 'x_train', 'y_train', 'x_tr', 'y_tr', '_trn', 'trn_', 'train')
         if isinstance(node, ast.Name):
             return any(m in node.id.lower() for m in train_markers)
         if isinstance(node, ast.Subscript):
-            return self._is_train_data(node.value)
+            return self._is_train_data(node.value, ctx)
         if isinstance(node, ast.Attribute):
-            return self._is_train_data(node.value)
+            return self._is_train_data(node.value, ctx)
         return False
 
     def check_node(self, node, ctx):
@@ -334,10 +347,17 @@ class T010_FitBeforeSplit(Rule):
                             'parser', 'formatter', 'logger', 'handler'}
             if obj_name and obj_name.lower() in safe_objects:
                 return []
+
+            # If arguments are contaminated, Flow engine F002 emits specific data contamination diagnostics
+            if ctx is not None and hasattr(ctx, 'test_contaminated_vars'):
+                for a in node.args:
+                    v_name = self._get_obj_name(a)
+                    if v_name and v_name.split('.')[0] in ctx.test_contaminated_vars:
+                        return []
             
             # Suppress if all positional arguments are clearly train-only variables
             # (e.g. X_train, y_train, train_df[features], X_tr)
-            if node.args and all(self._is_train_data(a) for a in node.args):
+            if node.args and all(self._is_train_data(a, ctx) for a in node.args):
                 return []
 
             if node.func.attr == 'fit':
@@ -524,6 +544,8 @@ class T017_TargetLaundering(Rule):
     def check_node(self, node, ctx):
         target_col_names = {'target', 'success', 'outcome', 'label', 'labels', 'y',
                             'profit', 'pnl', 'win', 'return', 'returns', 'future_price'}
+        train_markers = ('_train', 'train_', 'x_train', 'y_train', 'x_tr', 'y_tr', '_trn', 'trn_', 'train')
+        test_markers = ('_test', 'test_', 'x_test', 'y_test', 'x_val', 'y_val', '_val', 'val_', '_eval', 'eval_', 'holdout', 'validation')
         
         # 1. Detect target encoding / aggregation: df.groupby(...)['target'].transform('mean')
         if isinstance(node, ast.Call):
@@ -547,10 +569,23 @@ class T017_TargetLaundering(Rule):
                 
                 if col_name and any(t in col_name for t in target_col_names):
                     if isinstance(groupby_call, ast.Call) and isinstance(groupby_call.func, ast.Attribute) and groupby_call.func.attr == 'groupby':
+                        # Check the caller of groupby:
+                        caller_obj = groupby_call.func.value
+                        caller_name = ""
+                        if isinstance(caller_obj, ast.Name):
+                            caller_name = caller_obj.id.lower()
+                        elif isinstance(caller_obj, ast.Attribute):
+                            caller_name = caller_obj.attr.lower()
+                        
+                        # If the groupby is computed strictly within a training fold,
+                        # this is a VALID fold-local transformation! Do NOT flag it.
+                        if any(m in caller_name for m in train_markers):
+                            return []
+
                         return [Diagnostic(
                             node.lineno, node.col_offset, self.severity, self.code,
-                            f"Target laundering: grouping by target/outcome column '{col_name}' leaks future labels into features",
-                            "Compute regime/group statistics strictly on training data after splitting")]
+                            f"Target laundering: grouping by target/outcome column '{col_name}' across the un-split dataset leaks future test labels into features",
+                            "Compute regime/group target statistics strictly within the training fold (e.g. train_df.groupby()), and map to test fold with .fillna() for unseen regimes")]
 
         # 2. Detect assignment creating laundered columns: df['regime_success_rate'] = ...
         if isinstance(node, ast.Assign):
@@ -558,6 +593,33 @@ class T017_TargetLaundering(Rule):
                 if isinstance(target, ast.Subscript) and isinstance(target.slice, ast.Constant) and isinstance(target.slice.value, str):
                     col = target.slice.value.lower()
                     if any(marker in col for marker in ('success_rate', 'target_enc', 'target_mean', 'regime_success')):
+                        target_df = target.value
+                        target_df_name = ""
+                        if isinstance(target_df, ast.Name):
+                            target_df_name = target_df.id.lower()
+                        elif isinstance(target_df, ast.Attribute):
+                            target_df_name = target_df.attr.lower()
+
+                        # Case A: Assigning to a training fold dataframe -> valid fold-local transformation!
+                        if any(m in target_df_name for m in train_markers):
+                            return []
+
+                        # Case B: Assigning to test/inference fold dataframe (e.g. test_df['regime_success_rate'] = ...)
+                        if any(m in target_df_name for m in test_markers):
+                            # Check if unseen regimes at inference time are handled via .fillna()
+                            has_fillna = False
+                            if isinstance(node.value, ast.Call):
+                                if isinstance(node.value.func, ast.Attribute) and node.value.func.attr == 'fillna':
+                                    has_fillna = True
+                            
+                            if not has_fillna:
+                                return [Diagnostic(
+                                    node.lineno, node.col_offset, "warning", self.code,
+                                    f"Mapping target statistics to test data '{target_df_name}' without .fillna(): unseen categories/regimes at inference time will produce NaNs",
+                                    "Add .fillna(train_prior_mean) to safely handle unseen categories or market regimes at inference time")]
+                            return []
+
+                        # Case C: Assigning to un-split / full dataset dataframe (df['regime_success_rate'] = ...)
                         return [Diagnostic(
                             node.lineno, node.col_offset, self.severity, self.code,
                             f"Target laundering: column '{target.slice.value}' encodes target outcomes across full dataset",
@@ -595,19 +657,35 @@ class C002_TargetInFeatures(Rule):
     def check_node(self, node, ctx):
         if not isinstance(node, ast.Subscript):
             return []
+        train_markers = ('_train', 'train_', 'x_train', 'y_train', 'x_tr', 'y_tr', '_trn', 'trn_', 'train')
+        caller_name = ""
+        if isinstance(node.value, ast.Name):
+            caller_name = node.value.id.lower()
+        elif isinstance(node.value, ast.Attribute):
+            caller_name = node.value.attr.lower()
+
         # Detect df[['feature1', 'target']] where target-like names are used as features
         if isinstance(node.slice, ast.List):
-            target_names = {'target', 'y', 'label', 'labels', 'outcome', 'y_true',
-                            'future_price', 'future_return', 'next_close', 'next_price',
-                            'regime_success_rate', 'success_rate'}
+            raw_target_names = {'target', 'y', 'label', 'labels', 'outcome', 'y_true',
+                                'future_price', 'future_return', 'next_close', 'next_price'}
+            laundered_target_names = {'regime_success_rate', 'success_rate'}
             for elt in node.slice.elts:
                 if isinstance(elt, ast.Constant) and isinstance(elt.value, str):
                     val = elt.value.lower()
-                    if val in target_names or val.endswith('_success_rate') or val.endswith('_target_enc'):
+                    if val in raw_target_names:
+                        return [Diagnostic(
+                            node.lineno, node.col_offset, self.severity, self.code,
+                            f"Column '{elt.value}' appears to be a target variable used as a feature",
+                            "Remove target/outcome columns from feature inputs to avoid data leakage")]
+                    if val in laundered_target_names or val.endswith('_success_rate') or val.endswith('_target_enc'):
+                        # If caller is training data (e.g. train_df[['volatility', 'regime_success_rate']]),
+                        # this is a valid fold-local target encoding feature! Do NOT flag it.
+                        if any(m in caller_name for m in train_markers):
+                            continue
                         return [Diagnostic(
                             node.lineno, node.col_offset, self.severity, self.code,
-                            f"Column '{elt.value}' appears to be a target or target-laundered variable used as a feature",
-                            "Remove target/outcome and target-laundered columns from feature inputs to avoid data leakage")]
+                            f"Column '{elt.value}' appears to be a target or target-laundered variable used as a feature across full dataset",
+                            "Compute target-encoded features strictly fold-locally on training data")]
         return []
 
 

{pycheck_neuron-1.0.0 → pycheck_neuron-1.0.1/pycheck_neuron.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pycheck-neuron
-Version: 1.0.0
+Version: 1.0.1
 Summary: 32-rule static AST analyzer catching temporal lookahead bias, causal confusion, and uncertainty bugs in Python ML code. Zero dependencies.
 Author: NEURON Team
 Author-email: Fayo Ibrahim <fayo@neuron-lab.org>

{pycheck_neuron-1.0.0 → pycheck_neuron-1.0.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pycheck-neuron"
-version = "1.0.0"
+version = "1.0.1"
 description = "32-rule static AST analyzer catching temporal lookahead bias, causal confusion, and uncertainty bugs in Python ML code. Zero dependencies."
 readme = "README.md"
 authors = [

{pycheck_neuron-1.0.0 → pycheck_neuron-1.0.1}/tests/test_rules.py

@@ -515,4 +515,139 @@ X = df[['volatility', 'regime_success_rate']]
         assert any(d.code == "T017" for d in diags), "Target groupby aggregation must trigger T017"
         assert any(d.code == "C002" for d in diags), "regime_success_rate feature selection must trigger C002"
 
+    def test_provenance_concat_contamination(self):
+        code = """
+import pandas as pd
+from sklearn.preprocessing import StandardScaler
+
+X_train = pd.DataFrame({'a': [1, 2]})
+X_test = pd.DataFrame({'a': [3, 4]})
+
+X_train = pd.concat([X_train, X_test])
+scaler = StandardScaler()
+scaler.fit_transform(X_train)
+"""
+        tree = ast.parse(code)
+        tracker = TaintTracker()
+        results = tracker.analyze(tree, code.split('\n'))
+        assert any(d.code == "F002" for d in results), "Reassigned X_train with concatenated test data must trigger F002"
+        f002 = [d for d in results if d.code == "F002"][0]
+        assert "X_train" in f002.message
+        assert "test partition" in f002.message
+
+    def test_provenance_alias_and_reassignment(self):
+        code = """
+import pandas as pd
+from sklearn.preprocessing import StandardScaler
+
+X_eval = X_test
+X_train = pd.concat([X_train, X_eval])
+scaler = StandardScaler()
+scaler.fit(X_train)
+"""
+        tree = ast.parse(code)
+        tracker = TaintTracker()
+        results = tracker.analyze(tree, code.split('\n'))
+        assert any(d.code == "F002" for d in results), "Alias of test data in concat must trigger F002"
+
+    def test_provenance_helper_function(self):
+        code_leaky = """
+import pandas as pd
+from sklearn.preprocessing import StandardScaler
+
+def combine_folds(tr, te):
+    return pd.concat([tr, te])
+
+X_train = combine_folds(X_train, X_test)
+scaler = StandardScaler()
+scaler.fit_transform(X_train)
+"""
+        tree_leaky = ast.parse(code_leaky)
+        tracker = TaintTracker()
+        results_leaky = tracker.analyze(tree_leaky, code_leaky.split('\n'))
+        assert any(d.code == "F002" for d in results_leaky), "Helper function combining test data must trigger F002"
+
+        code_clean = """
+from sklearn.preprocessing import StandardScaler
+
+def clean_data(data):
+    return data.dropna()
+
+clean_train = clean_data(X_train)
+scaler = StandardScaler()
+scaler.fit_transform(clean_train)
+"""
+        tree_clean = ast.parse(code_clean)
+        tracker_clean = TaintTracker()
+        results_clean = tracker_clean.analyze(tree_clean, code_clean.split('\n'))
+        assert not any(d.code == "F002" for d in results_clean), "Clean helper function on training data must not trigger F002"
+
+    def test_t017_fold_local_target_encoding_clean(self):
+        code = """
+# Legitimate fold-local target encoding
+regime_stats = train_df.groupby('regime')['target'].mean()
+global_prior = train_df['target'].mean()
+
+train_df['regime_success_rate'] = train_df['regime'].map(regime_stats)
+test_df['regime_success_rate'] = test_df['regime'].map(regime_stats).fillna(global_prior)
+X_train = train_df[['volatility', 'regime_success_rate']]
+"""
+        tree = ast.parse(code)
+        from pycheck.rules import T017_TargetLaundering, C002_TargetInFeatures
+        t017 = T017_TargetLaundering()
+        c002 = C002_TargetInFeatures()
+        diags = []
+        for node in ast.walk(tree):
+            diags.extend(t017.check_node(node, None))
+            diags.extend(c002.check_node(node, None))
+        assert not any(d.code == "T017" for d in diags), f"Fold-local target encoding must not trigger T017: {diags}"
+        assert not any(d.code == "C002" for d in diags), f"Fold-local target feature on train_df must not trigger C002: {diags}"
+
+    def test_t017_global_target_encoding_flagged(self):
+        code = """
+# Global target laundering across un-split dataset
+df['regime_success_rate'] = df.groupby('regime')['target'].transform('mean')
+"""
+        tree = ast.parse(code)
+        from pycheck.rules import T017_TargetLaundering
+        t017 = T017_TargetLaundering()
+        diags = []
+        for node in ast.walk(tree):
+            diags.extend(t017.check_node(node, None))
+        assert any(d.code == "T017" for d in diags), "Global groupby transform on target must trigger T017"
+
+    def test_t017_unseen_regimes_inference_warning(self):
+        code = """
+regime_stats = train_df.groupby('regime')['target'].mean()
+# Missing .fillna() for unseen regimes at test inference time
+test_df['regime_success_rate'] = test_df['regime'].map(regime_stats)
+"""
+        tree = ast.parse(code)
+        from pycheck.rules import T017_TargetLaundering
+        t017 = T017_TargetLaundering()
+        diags = []
+        for node in ast.walk(tree):
+            diags.extend(t017.check_node(node, None))
+        assert any(d.code == "T017" and d.severity == "warning" for d in diags), "Mapping to test_df without fillna must trigger unseen regimes warning"
+
+    def test_provenance_full_pipeline_analyze_file(self, tmp_path):
+        test_file = tmp_path / "leaky_provenance.py"
+        test_file.write_text("""
+import pandas as pd
+from sklearn.preprocessing import StandardScaler
+
+X_train = pd.DataFrame({'a': [1, 2]})
+X_test = pd.DataFrame({'a': [3, 4]})
+
+X_train = pd.concat([X_train, X_test])
+scaler = StandardScaler()
+scaler.fit_transform(X_train)
+""", encoding='utf-8')
+        diagnostics, _ = analyze_file(str(test_file))
+        assert any(d['code'] == "F002" for d in diagnostics), f"Full analyze_file must detect F002, got: {diagnostics}"
+        f002 = [d for d in diagnostics if d['code'] == "F002"][0]
+        assert "X_train" in f002['message']
+        assert "test partition" in f002['message']
+
+
 

pycheck_neuron-1.0.0/pycheck/flow.py

@@ -1,265 +0,0 @@
-"""
-PyCheck Data Flow Engine — Taint Propagation Analysis
-Tracks how data flows through assignments, function returns, and method chains
-to detect when future-tainted data reaches training/prediction sinks.
-"""
-
-import ast
-from typing import Dict, Set, List, Optional, Tuple
-from pycheck.rules import Diagnostic
-
-
-class TaintSource:
-    """Represents a source of tainted (future-leaked) data."""
-    def __init__(self, var_name: str, reason: str, line: int, col: int):
-        self.var_name = var_name
-        self.reason = reason
-        self.line = line
-        self.col = col
-
-
-class TaintTracker:
-    """
-    Tracks taint propagation through Python AST.
-    
-    Taint sources: shift(-n), bfill(), future column access, etc.
-    Taint sinks: .fit(), .predict(), model training calls.
-    
-    When tainted data flows into a sink, a diagnostic is emitted.
-    """
-    
-    def __init__(self):
-        self.tainted_vars: Dict[str, TaintSource] = {}
-        self.diagnostics: List[Diagnostic] = []
-        self.constants: Dict[str, any] = {}
-    
-    def analyze(self, tree: ast.AST, source_lines: List[str]) -> List[Diagnostic]:
-        """Run taint analysis on a full AST."""
-        self.tainted_vars.clear()
-        self.diagnostics.clear()
-        self.constants.clear()
-        
-        # Two-pass analysis:
-        # Pass 1: Identify all taint sources (assignments from leaky operations)
-        # Pass 2: Track propagation and detect sinks
-        self._collect_taints(tree)
-        self._check_sinks(tree)
-        
-        return self.diagnostics
-    
-    def _extract_constant(self, expr):
-        if isinstance(expr, ast.Constant) and isinstance(expr.value, (int, float)):
-            return expr.value
-        if isinstance(expr, ast.UnaryOp) and isinstance(expr.op, ast.USub):
-            if isinstance(expr.operand, ast.Constant) and isinstance(expr.operand.value, (int, float)):
-                return -expr.operand.value
-        return None
-
-    def _is_negative_period(self, expr) -> bool:
-        if isinstance(expr, ast.UnaryOp) and isinstance(expr.op, ast.USub):
-            return True
-        if isinstance(expr, ast.Constant) and isinstance(expr.value, (int, float)) and expr.value < 0:
-            return True
-        if isinstance(expr, ast.Name):
-            val = self.constants.get(expr.id)
-            if val is not None and isinstance(val, (int, float)) and val < 0:
-                return True
-        return False
-
-    def _collect_taints(self, tree: ast.AST):
-        """Pass 1: Walk the AST and mark variables that receive tainted data."""
-        for node in ast.walk(tree):
-            if isinstance(node, ast.Assign):
-                for target in node.targets:
-                    if isinstance(target, ast.Name):
-                        val = self._extract_constant(node.value)
-                        if val is not None:
-                            self.constants[target.id] = val
-
-        for node in ast.walk(tree):
-            if isinstance(node, ast.Assign):
-                self._check_assignment_taint(node)
-            elif isinstance(node, ast.AugAssign):
-                self._check_augassign_taint(node)
-    
-    def _check_assignment_taint(self, node: ast.Assign):
-        """Check if an assignment creates or propagates a taint."""
-        # Get the taint from the RHS
-        taint_reason = self._expr_taint(node.value)
-        if taint_reason is None:
-            return
-        
-        # Mark all LHS targets as tainted
-        for target in node.targets:
-            var_name = self._get_var_name(target)
-            if var_name:
-                self.tainted_vars[var_name] = TaintSource(
-                    var_name, taint_reason, node.lineno, node.col_offset)
-    
-    def _check_augassign_taint(self, node: ast.AugAssign):
-        """Check if an augmented assignment (+=, etc.) propagates taint."""
-        taint_reason = self._expr_taint(node.value)
-        if taint_reason:
-            var_name = self._get_var_name(node.target)
-            if var_name:
-                self.tainted_vars[var_name] = TaintSource(
-                    var_name, taint_reason, node.lineno, node.col_offset)
-    
-    def _expr_taint(self, node: ast.AST) -> Optional[str]:
-        """
-        Determine if an expression produces tainted data.
-        Returns the taint reason string, or None if clean.
-        """
-        if isinstance(node, ast.Call):
-            return self._call_taint(node)
-        
-        if isinstance(node, ast.Name):
-            # Propagation: if RHS is a tainted variable, the assignment is tainted
-            if node.id in self.tainted_vars:
-                return f"derived from tainted variable '{node.id}' ({self.tainted_vars[node.id].reason})"
-        
-        if isinstance(node, ast.BinOp):
-            # Binary operations propagate taint from either operand
-            left_taint = self._expr_taint(node.left)
-            right_taint = self._expr_taint(node.right)
-            return left_taint or right_taint
-        
-        if isinstance(node, ast.Subscript):
-            # Subscript propagates taint from the value being indexed
-            val_taint = self._expr_taint(node.value)
-            if val_taint:
-                return val_taint
-            # Check if indexing a tainted column name
-            if isinstance(node.slice, ast.Constant) and isinstance(node.slice.value, str):
-                col_lower = node.slice.value.lower()
-                # Only flag column names that unambiguously refer to future data.
-                # 'target' and 'label' are intentionally excluded — they are standard
-                # ML label column names and produce false positives when accessed on a
-                # correctly split train/test dataset.
-                future_cols = {'future_price', 'future_return', 'next_close', 'next_price',
-                               'next_day', 'tomorrow', 'price_change_next_day', 'future_avg_price'}
-                if any(f in col_lower for f in future_cols):
-                    return f"accessing future-named column '{node.slice.value}'"
-        
-        if isinstance(node, ast.Attribute):
-            return self._expr_taint(node.value)
-        
-        return None
-    
-    def _call_taint(self, node: ast.Call) -> Optional[str]:
-        """Check if a function call produces tainted output."""
-        if isinstance(node.func, ast.Attribute):
-            method = node.func.attr
-            
-            # Direct taint sources
-            if method == 'shift':
-                for arg in node.args:
-                    if self._is_negative_period(arg):
-                        return "shift() with negative period (future access)"
-                for kw in node.keywords:
-                    if kw.arg == 'periods' and self._is_negative_period(kw.value):
-                        return "shift() with negative period (future access)"
-            
-            if method in ('bfill', 'backfill'):
-                return f".{method}() backfills from future values"
-            
-            if method == 'fillna':
-                for kw in node.keywords:
-                    if kw.arg == 'method' and isinstance(kw.value, ast.Constant):
-                        if kw.value.value in ('bfill', 'backfill'):
-                            return f".fillna(method='{kw.value.value}') backfills from future"
-            
-            if method == 'diff':
-                for arg in node.args:
-                    if self._is_negative_period(arg):
-                        return ".diff() with negative period (future access)"
-                for kw in node.keywords:
-                    if kw.arg == 'periods' and self._is_negative_period(kw.value):
-                        return ".diff() with negative period (future access)"
-            
-            # Taint propagation through method chains
-            obj_taint = self._expr_taint(node.func.value)
-            if obj_taint:
-                # Most method calls on tainted data produce tainted output
-                non_taint_methods = {'shape', 'dtype', 'dtypes', 'columns', 'index',
-                                     'describe', 'info', 'head', 'tail', 'len', '__len__'}
-                if method not in non_taint_methods:
-                    return obj_taint
-        
-        if isinstance(node.func, ast.Name):
-            # Check if calling a function with tainted arguments
-            for arg in node.args:
-                arg_taint = self._expr_taint(arg)
-                if arg_taint:
-                    return f"function receives tainted input: {arg_taint}"
-        
-        # R2-9 fix: Detect getattr(obj, 'shift')(-1) and similar patterns
-        # Pattern: call(getattr(obj, 'shift'), -1) => obj.shift(-1)
-        if isinstance(node.func, ast.Call):
-            inner_call = node.func
-            if isinstance(inner_call.func, ast.Name) and inner_call.func.id == 'getattr':
-                if len(inner_call.args) >= 2 and isinstance(inner_call.args[1], ast.Constant):
-                    attr_name = inner_call.args[1].value
-                    if attr_name == 'shift':
-                        for arg in node.args:
-                            if self._is_negative_period(arg):
-                                return "getattr(..., 'shift')() with negative period (future access)"
-                        for kw in node.keywords:
-                            if kw.arg == 'periods' and self._is_negative_period(kw.value):
-                                return "getattr(..., 'shift')() with negative period (future access)"
-                    elif attr_name == 'diff':
-                        for arg in node.args:
-                            if self._is_negative_period(arg):
-                                return "getattr(..., 'diff')() with negative period (future access)"
-                        for kw in node.keywords:
-                            if kw.arg == 'periods' and self._is_negative_period(kw.value):
-                                return "getattr(..., 'diff')() with negative period (future access)"
-                    elif attr_name in ('bfill', 'backfill'):
-                        return f"getattr(..., '{attr_name}')() backfills from future values"
-        
-        return None
-    
-    def _check_sinks(self, tree: ast.AST):
-        """Pass 2: Check if tainted data flows into sensitive sinks."""
-        for node in ast.walk(tree):
-            if not isinstance(node, ast.Call):
-                continue
-            
-            if isinstance(node.func, ast.Attribute):
-                method = node.func.attr
-                sinks = {'fit', 'fit_transform', 'fit_predict', 'train',
-                          'partial_fit', 'predict', 'score'}
-                
-                if method in sinks:
-                    # Check if any positional argument is tainted
-                    for arg in node.args:
-                        arg_taint = self._expr_taint(arg)
-                        if arg_taint:
-                            self.diagnostics.append(Diagnostic(
-                                node.lineno, node.col_offset, "error", "F001",
-                                f"Tainted data flows into .{method}(): {arg_taint}",
-                                "Ensure training/prediction data is free of future information"))
-                    # R2-8 fix: Check keyword arguments for taint as well
-                    for kw in node.keywords:
-                        kw_taint = self._expr_taint(kw.value)
-                        if kw_taint:
-                            kw_name = kw.arg if kw.arg else "**kwargs"
-                            self.diagnostics.append(Diagnostic(
-                                node.lineno, node.col_offset, "error", "F001",
-                                f"Tainted data flows into .{method}() via keyword '{kw_name}': {kw_taint}",
-                                "Ensure training/prediction data is free of future information"))
-    
-    def _get_var_name(self, node: ast.AST) -> Optional[str]:
-        """Extract variable name from an assignment target."""
-        if isinstance(node, ast.Name):
-            return node.id
-        if isinstance(node, ast.Subscript):
-            # df['col_name'] — return as df.col_name
-            base = self._get_var_name(node.value)
-            if isinstance(node.slice, ast.Constant) and isinstance(node.slice.value, str):
-                return f"{base}.{node.slice.value}" if base else node.slice.value
-            return base
-        if isinstance(node, ast.Attribute):
-            base = self._get_var_name(node.value)
-            return f"{base}.{node.attr}" if base else node.attr
-        return None