code/flash-linear-attention/scripts/find_dependent_tests.py

import ast
import os
import sys
from collections import defaultdict
from functools import cache
from pathlib import Path

DEBUG_MODE = os.environ.get("DEBUG_MODE", "False").lower() in ("true", "1", "yes")
DEBUG_TEST_FILE = os.environ.get("DEBUG_TEST_FILE", "NULL").lower()


@cache
def parse_file(file_path):
    try:
        with open(file_path, encoding="utf-8") as f:
            return ast.parse(f.read(), filename=file_path)
    except (SyntaxError, FileNotFoundError, UnicodeDecodeError):
        return None


def get_definitions_from_tree(tree) -> set:
    if not tree:
        return set()
    definitions = set()
    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            definitions.add(node.name)
    return definitions


def get_imports_from_tree(tree) -> set:
    if not tree:
        return set()
    imports = set()
    for node in tree.body:
        if isinstance(node, ast.ImportFrom):
            for alias in node.names:
                imports.add(alias.asname or alias.name)
        elif isinstance(node, ast.Import):
            for alias in node.names:
                imports.add(alias.asname or alias.name.split('.')[0])
    return imports


class DependencyFinder:
    def __init__(self, search_dirs, test_dir):
        self.test_dir = Path(test_dir).resolve()
        models_test_dir = self.test_dir / "models"

        source_files = [p for s_dir in search_dirs for p in Path(s_dir).resolve().rglob("*.py") if p.name != '__init__.py']
        test_scope = os.environ.get("TEST_SCOPE", "ALL").upper()
        if test_scope == "MODELS_ONLY":
            test_files = [p for p in models_test_dir.rglob("*.py") if p.name != '__init__.py']
        elif test_scope == "EXCLUDE_MODELS":
            all_files = self.test_dir.rglob("*.py")
            test_files = [p for p in all_files if p.name != '__init__.py' and models_test_dir not in p.parents]
        else:
            test_files = [p for p in self.test_dir.rglob("*.py") if p.name != '__init__.py']
        self.all_project_files = source_files + test_files
        self.all_test_files = set(test_files)

        self.file_to_definitions = {}
        self.file_to_imports = {}
        self.symbol_to_file_map = defaultdict(set)
        for file_path in self.all_project_files:
            tree = parse_file(file_path)
            definitions = get_definitions_from_tree(tree)
            imports = get_imports_from_tree(tree)
            self.file_to_definitions[file_path] = definitions
            self.file_to_imports[file_path] = imports
            for defn in definitions:
                self.symbol_to_file_map[defn].add(file_path)

    def _print_dependency_chain(self, symbol, symbol_chain_map):
        chain = []
        current_symbol = symbol
        while current_symbol is not None:
            # Find the file where the symbol is defined
            # In case of multiple definitions, we take the first one found
            file_path = next(iter(self.symbol_to_file_map.get(current_symbol, ["Unknown File"])), "Unknown File")
            chain.append(f"{current_symbol} @ {file_path}")
            current_symbol = symbol_chain_map.get(current_symbol)

        chain.reverse()
        print("  - Dependency Chain:", " -> ".join(chain), file=sys.stderr)

    def find_dependent_tests(self, changed_files_str: list, max_depth=4) -> set:
        changed_files = {Path(f).resolve() for f in changed_files_str}

        initial_configs_to_add = set()
        for file in changed_files:
            if 'modeling_' in file.stem and 'models' in str(file):
                model_name = file.stem.replace('modeling_', '')
                config_file = file.parent / f"configuration_{model_name}.py"
                if config_file.is_file():
                    initial_configs_to_add.add(config_file)
        changed_files.update(initial_configs_to_add)

        symbol_chain_map = {}
        all_affected_symbols = set()
        symbols_to_trace = set()

        for file_path in changed_files:
            if file_path.name == '__init__.py':
                continue
            new_defs = self.file_to_definitions.get(file_path, set())
            symbols_to_trace.update(new_defs)
            all_affected_symbols.update(new_defs)
            for defn in new_defs:
                symbol_chain_map[defn] = None

        for i in range(max_depth):
            if not symbols_to_trace:
                break

            next_layer_files = set()

            # Find files that import the current symbols to trace
            # And for each new definition, link it to the symbol that triggered it
            newly_added_definitions = set()
            for file_path, imported_symbols in self.file_to_imports.items():
                triggers = symbols_to_trace.intersection(imported_symbols)
                if triggers:
                    next_layer_files.add(file_path)
                    defs_in_file = self.file_to_definitions.get(file_path, set())
                    # For simplicity, we link all new definitions in this file to the first trigger found
                    first_trigger = next(iter(triggers))
                    for defn in defs_in_file:
                        if defn not in all_affected_symbols:
                            symbol_chain_map[defn] = first_trigger
                            newly_added_definitions.add(defn)

            # This heuristic is now also applied at each dependency level
            config_files_to_add = set()
            for file in next_layer_files:
                if 'modeling_' in file.stem and 'models' in str(file):
                    model_name = file.stem.replace('modeling_', '')
                    config_file = file.parent / f"configuration_{model_name}.py"
                    if config_file.is_file():
                        config_files_to_add.add(config_file)

            # For config files, we don't have a clear trigger, so we can't map their chain
            for config_file in config_files_to_add:
                defs_in_file = self.file_to_definitions.get(config_file, set())
                for defn in defs_in_file:
                    if defn not in all_affected_symbols:
                        symbol_chain_map[defn] = "CONFIG_HEURISTIC"  # Special marker
                        newly_added_definitions.add(defn)

            next_layer_files.update(config_files_to_add)

            symbols_to_trace = newly_added_definitions
            all_affected_symbols.update(symbols_to_trace)

        dependent_tests = set()

        affected_source_file_stems = set()
        for s in all_affected_symbols:
            if s in self.symbol_to_file_map:
                for file_path in self.symbol_to_file_map[s]:
                    affected_source_file_stems.add(file_path.stem)

        for test_file in self.all_test_files:
            imported_in_test = self.file_to_imports.get(test_file, set())

            if not all_affected_symbols.isdisjoint(imported_in_test):
                if DEBUG_MODE and DEBUG_TEST_FILE in str(test_file).lower():
                    imported_symbols = [s for s in imported_in_test if s in all_affected_symbols]
                    print(
                        f"DEBUG: Test file {test_file} is included because it imports affected symbols: {imported_symbols}", file=sys.stderr)  # noqa: E501
                    for symbol in imported_symbols:
                        self._print_dependency_chain(symbol, symbol_chain_map)
                dependent_tests.add(str(test_file))
                continue

            if not affected_source_file_stems.isdisjoint(imported_in_test):
                if DEBUG_MODE and DEBUG_TEST_FILE in str(test_file).lower():
                    imported_files = [f for f in imported_in_test if f in affected_source_file_stems]
                    print(
                        f"DEBUG: Test file {test_file} is included because it imports a symbol matching an affected file stem: {imported_files}", file=sys.stderr)  # noqa: E501
                dependent_tests.add(str(test_file))
        for changed_file in changed_files:
            if changed_file in self.all_test_files:
                dependent_tests.add(str(changed_file))

        return dependent_tests


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python find_dependent_tests.py <file1> <file2> ...")
        sys.exit(1)

    all_args_string = " ".join(sys.argv[1:])
    changed_files = all_args_string.split()

    BLACKLIST = ['fla/utils.py', 'utils/convert_from_llama.py', 'utils/convert_from_rwkv6.py', 'utils/convert_from_rwkv7.py']
    changed_files = [file for file in changed_files if not any(file.endswith(b) for b in BLACKLIST)]

    changed_files = [file for file in changed_files if file.endswith('.py')]

    current_dir = Path(__file__).parent.resolve()
    test_dir = current_dir.parent / "tests"
    search_dir = current_dir.parent / "fla"

    finder = DependencyFinder(search_dirs=[search_dir], test_dir=test_dir)
    dependent_tests = finder.find_dependent_tests(changed_files)

    if dependent_tests:
        print(" ".join(sorted(list(dependent_tests))))