nogic 0.0.1__py3-none-any.whl

nogic/parsing/js_extractor.py

@@ -0,0 +1,674 @@
+"""
+JavaScript/TypeScript code extractor using Tree-sitter.
+
+Handles:
+- Function declarations
+- Arrow functions assigned to variables
+- Class declarations with methods
+- React components (functions returning JSX)
+- ES6 and CommonJS imports
+"""
+
+from tree_sitter import Node
+
+from .types import (
+    ExtractedFunction,
+    ExtractedClass,
+    ExtractedCall,
+    ExtractedImport,
+)
+
+
+class JavaScriptExtractor:
+    """Extractor for JavaScript and TypeScript code."""
+
+    def __init__(self, source: bytes, module_qn: str, is_typescript: bool = False):
+        self.source = source
+        self.module_qn = module_qn
+        self.is_typescript = is_typescript
+
+    def get_node_text(self, node: Node) -> str:
+        """Get text content of a node."""
+        return self.source[node.start_byte : node.end_byte].decode("utf-8")
+
+    def get_node_lines(self, node: Node) -> tuple[int, int]:
+        """Get start and end line numbers (1-indexed)."""
+        return node.start_point[0] + 1, node.end_point[0] + 1
+
+    def extract_functions(self, root: Node) -> list[ExtractedFunction]:
+        """Extract all top-level functions (not methods)."""
+        functions: list[ExtractedFunction] = []
+        self._extract_functions_recursive(root, functions)
+        return functions
+
+    def _extract_functions_recursive(
+        self, node: Node, functions: list[ExtractedFunction]
+    ) -> None:
+        """Recursively extract functions."""
+        # Function declarations
+        if node.type == "function_declaration":
+            if not self._is_inside_class(node):
+                func = self._extract_function_declaration(node)
+                if func:
+                    functions.append(func)
+
+        # Arrow functions assigned to variables: const foo = () => {}
+        elif node.type in ("lexical_declaration", "variable_declaration"):
+            for declarator in self._get_declarators(node):
+                func = self._extract_arrow_function(declarator)
+                if func:
+                    functions.append(func)
+
+        # Export statements may contain functions
+        elif node.type in ("export_statement", "export_default_declaration"):
+            for child in node.children:
+                if child.type == "function_declaration":
+                    func = self._extract_function_declaration(child)
+                    if func:
+                        functions.append(func)
+                elif child.type in ("lexical_declaration", "variable_declaration"):
+                    for declarator in self._get_declarators(child):
+                        func = self._extract_arrow_function(declarator)
+                        if func:
+                            functions.append(func)
+
+        # Recurse but skip class bodies
+        if node.type not in ("class_body", "class_declaration", "class_expression"):
+            for child in node.children:
+                self._extract_functions_recursive(child, functions)
+
+    def _get_declarators(self, decl_node: Node) -> list[Node]:
+        """Get variable_declarator nodes from a declaration."""
+        return [c for c in decl_node.children if c.type == "variable_declarator"]
+
+    def _is_inside_class(self, node: Node) -> bool:
+        """Check if node is inside a class body."""
+        parent = node.parent
+        while parent:
+            if parent.type == "class_body":
+                return True
+            parent = parent.parent
+        return False
+
+    def _extract_function_declaration(self, node: Node) -> ExtractedFunction | None:
+        """Extract from function_declaration node."""
+        name_node = node.child_by_field_name("name")
+        if not name_node:
+            return None
+
+        name = self.get_node_text(name_node)
+        qualified_name = f"{self.module_qn}.{name}"
+        start_line, end_line = self.get_node_lines(node)
+
+        # Get parameters
+        parameters = self._extract_parameters(node)
+        params_text = self._get_parameters_text(node)
+
+        # Check if async
+        is_async = self._is_async(node)
+
+        # Build signature
+        prefix = "async " if is_async else ""
+        signature = f"{prefix}function {name}{params_text}"
+
+        # Get JSDoc comment as docstring
+        docstring = self._extract_jsdoc(node)
+
+        return ExtractedFunction(
+            name=name,
+            qualified_name=qualified_name,
+            start_line=start_line,
+            end_line=end_line,
+            signature=signature,
+            docstring=docstring,
+            decorators=[],
+            is_async=is_async,
+            is_method=False,
+            parameters=parameters,
+            source_code=self.get_node_text(node),
+            node=node,
+        )
+
+    def _extract_arrow_function(self, declarator: Node) -> ExtractedFunction | None:
+        """Extract arrow function from variable_declarator."""
+        name_node = declarator.child_by_field_name("name")
+        value_node = declarator.child_by_field_name("value")
+
+        if not name_node or not value_node:
+            return None
+
+        # Check if value is an arrow function or function expression
+        if value_node.type not in (
+            "arrow_function",
+            "function_expression",
+            "function",
+        ):
+            return None
+
+        name = self.get_node_text(name_node)
+        qualified_name = f"{self.module_qn}.{name}"
+
+        # Use the whole declaration for line numbers
+        parent = declarator.parent
+        if parent:
+            start_line, end_line = self.get_node_lines(parent)
+        else:
+            start_line, end_line = self.get_node_lines(declarator)
+
+        # Get parameters from the arrow/function
+        parameters = self._extract_parameters(value_node)
+        params_text = self._get_parameters_text(value_node)
+
+        # Check if async
+        is_async = self._is_async(value_node)
+
+        # Build signature
+        prefix = "async " if is_async else ""
+        if value_node.type == "arrow_function":
+            signature = f"const {name} = {prefix}{params_text} =>"
+        else:
+            signature = f"const {name} = {prefix}function{params_text}"
+
+        # Get JSDoc
+        docstring = self._extract_jsdoc(parent) if parent else None
+
+        return ExtractedFunction(
+            name=name,
+            qualified_name=qualified_name,
+            start_line=start_line,
+            end_line=end_line,
+            signature=signature,
+            docstring=docstring,
+            decorators=[],
+            is_async=is_async,
+            is_method=False,
+            parameters=parameters,
+            source_code=self.get_node_text(parent)
+            if parent
+            else self.get_node_text(declarator),
+            node=value_node,
+        )
+
+    def _extract_parameters(self, func_node: Node) -> list[str]:
+        """Extract parameter names."""
+        parameters: list[str] = []
+        params_node = func_node.child_by_field_name("parameters")
+        if not params_node:
+            # Arrow functions might have params directly
+            params_node = func_node.child_by_field_name("parameter")
+            if params_node and params_node.type == "identifier":
+                return [self.get_node_text(params_node)]
+            return parameters
+
+        for child in params_node.children:
+            if child.type == "identifier":
+                parameters.append(self.get_node_text(child))
+            elif child.type in ("required_parameter", "optional_parameter"):
+                # TypeScript typed parameters
+                pattern = child.child_by_field_name("pattern")
+                if pattern and pattern.type == "identifier":
+                    parameters.append(self.get_node_text(pattern))
+            elif child.type == "rest_pattern":
+                # ...args
+                for c in child.children:
+                    if c.type == "identifier":
+                        parameters.append("..." + self.get_node_text(c))
+            elif child.type == "assignment_pattern":
+                # Default parameters: x = value
+                left = child.child_by_field_name("left")
+                if left and left.type == "identifier":
+                    parameters.append(self.get_node_text(left))
+
+        return parameters
+
+    def _get_parameters_text(self, func_node: Node) -> str:
+        """Get parameters as text."""
+        params_node = func_node.child_by_field_name("parameters")
+        if params_node:
+            return self.get_node_text(params_node)
+        # Single parameter arrow function without parens
+        param = func_node.child_by_field_name("parameter")
+        if param:
+            return f"({self.get_node_text(param)})"
+        return "()"
+
+    def _is_async(self, node: Node) -> bool:
+        """Check if function is async."""
+        # Check for 'async' keyword in children
+        for child in node.children:
+            if child.type == "async":
+                return True
+            if self.get_node_text(child) == "async":
+                return True
+        return False
+
+    def _extract_jsdoc(self, node: Node) -> str | None:
+        """Extract JSDoc comment preceding a node."""
+        # Look for comment in previous siblings
+        if node.prev_sibling and node.prev_sibling.type == "comment":
+            comment = self.get_node_text(node.prev_sibling)
+            if comment.startswith("/**"):
+                # Strip /** and */
+                return comment[3:-2].strip()
+        return None
+
+    def extract_classes(self, root: Node) -> list[ExtractedClass]:
+        """Extract all classes from the AST."""
+        classes: list[ExtractedClass] = []
+        self._extract_classes_recursive(root, classes)
+        return classes
+
+    def _extract_classes_recursive(
+        self, node: Node, classes: list[ExtractedClass]
+    ) -> None:
+        """Recursively extract classes."""
+        if node.type in ("class_declaration", "class_expression"):
+            cls = self._extract_class(node)
+            if cls:
+                classes.append(cls)
+
+        # Also check export statements
+        elif node.type in ("export_statement", "export_default_declaration"):
+            for child in node.children:
+                if child.type in ("class_declaration", "class_expression"):
+                    cls = self._extract_class(child)
+                    if cls:
+                        classes.append(cls)
+
+        for child in node.children:
+            if child.type not in ("class_body",):
+                self._extract_classes_recursive(child, classes)
+
+    def _extract_class(self, node: Node) -> ExtractedClass | None:
+        """Extract class information."""
+        name_node = node.child_by_field_name("name")
+        if not name_node:
+            # Anonymous class
+            return None
+
+        name = self.get_node_text(name_node)
+        qualified_name = f"{self.module_qn}.{name}"
+        start_line, end_line = self.get_node_lines(node)
+
+        # Get base classes (extends)
+        bases = self._extract_class_heritage(node)
+
+        # Get JSDoc
+        docstring = self._extract_jsdoc(node)
+
+        # Get decorators (TypeScript)
+        decorators = self._extract_decorators(node)
+
+        # Get methods
+        methods = self._extract_methods(node, name)
+
+        return ExtractedClass(
+            name=name,
+            qualified_name=qualified_name,
+            start_line=start_line,
+            end_line=end_line,
+            docstring=docstring,
+            decorators=decorators,
+            bases=bases,
+            methods=methods,
+            source_code=self.get_node_text(node),
+            node=node,
+        )
+
+    def _extract_class_heritage(self, class_node: Node) -> list[str]:
+        """Extract base class names."""
+        bases: list[str] = []
+        for child in class_node.children:
+            if child.type == "class_heritage":
+                for c in child.children:
+                    if c.type == "extends_clause":
+                        # Get the class being extended
+                        for cc in c.children:
+                            if cc.type == "identifier":
+                                bases.append(self.get_node_text(cc))
+                            elif cc.type == "member_expression":
+                                bases.append(self.get_node_text(cc))
+        return bases
+
+    def _extract_decorators(self, node: Node) -> list[str]:
+        """Extract decorators (TypeScript)."""
+        decorators: list[str] = []
+        for child in node.children:
+            if child.type == "decorator":
+                decorators.append(self.get_node_text(child))
+        return decorators
+
+    def _extract_methods(
+        self, class_node: Node, class_name: str
+    ) -> list[ExtractedFunction]:
+        """Extract methods from a class."""
+        methods: list[ExtractedFunction] = []
+        body_node = class_node.child_by_field_name("body")
+        if not body_node:
+            return methods
+
+        for child in body_node.children:
+            if child.type == "method_definition":
+                method = self._extract_method(child, class_name)
+                if method:
+                    methods.append(method)
+            elif child.type == "public_field_definition":
+                # Arrow function as class field: foo = () => {}
+                method = self._extract_field_method(child, class_name)
+                if method:
+                    methods.append(method)
+
+        return methods
+
+    def _extract_method(
+        self, node: Node, class_name: str
+    ) -> ExtractedFunction | None:
+        """Extract method from method_definition node."""
+        name_node = node.child_by_field_name("name")
+        if not name_node:
+            return None
+
+        name = self.get_node_text(name_node)
+        qualified_name = f"{self.module_qn}.{class_name}.{name}"
+        start_line, end_line = self.get_node_lines(node)
+
+        parameters = self._extract_parameters(node)
+        params_text = self._get_parameters_text(node)
+
+        is_async = self._is_async(node)
+
+        # Check for getter/setter
+        prefix = ""
+        for child in node.children:
+            if child.type == "get":
+                prefix = "get "
+            elif child.type == "set":
+                prefix = "set "
+            elif child.type == "static":
+                prefix = "static " + prefix
+
+        if is_async:
+            prefix = prefix + "async "
+
+        signature = f"{prefix}{name}{params_text}"
+
+        docstring = self._extract_jsdoc(node)
+        decorators = self._extract_decorators(node)
+
+        return ExtractedFunction(
+            name=name,
+            qualified_name=qualified_name,
+            start_line=start_line,
+            end_line=end_line,
+            signature=signature,
+            docstring=docstring,
+            decorators=decorators,
+            is_async=is_async,
+            is_method=True,
+            class_name=class_name,
+            parameters=parameters,
+            source_code=self.get_node_text(node),
+            node=node,
+        )
+
+    def _extract_field_method(
+        self, node: Node, class_name: str
+    ) -> ExtractedFunction | None:
+        """Extract arrow function from class field."""
+        name_node = node.child_by_field_name("name")
+        value_node = node.child_by_field_name("value")
+
+        if not name_node or not value_node:
+            return None
+
+        if value_node.type not in ("arrow_function", "function_expression"):
+            return None
+
+        name = self.get_node_text(name_node)
+        qualified_name = f"{self.module_qn}.{class_name}.{name}"
+        start_line, end_line = self.get_node_lines(node)
+
+        parameters = self._extract_parameters(value_node)
+        params_text = self._get_parameters_text(value_node)
+
+        is_async = self._is_async(value_node)
+
+        prefix = "async " if is_async else ""
+        signature = f"{name} = {prefix}{params_text} =>"
+
+        docstring = self._extract_jsdoc(node)
+
+        return ExtractedFunction(
+            name=name,
+            qualified_name=qualified_name,
+            start_line=start_line,
+            end_line=end_line,
+            signature=signature,
+            docstring=docstring,
+            decorators=[],
+            is_async=is_async,
+            is_method=True,
+            class_name=class_name,
+            parameters=parameters,
+            source_code=self.get_node_text(node),
+            node=value_node,
+        )
+
+    def extract_calls(self, node: Node, caller_qn: str) -> list[ExtractedCall]:
+        """Extract function calls from within a function/method node."""
+        calls: list[ExtractedCall] = []
+        self._extract_calls_recursive(node, caller_qn, calls)
+        return calls
+
+    def _extract_calls_recursive(
+        self, node: Node, caller_qn: str, calls: list[ExtractedCall]
+    ) -> None:
+        """Recursively extract call expressions."""
+        if node.type == "call_expression":
+            call = self._extract_single_call(node, caller_qn)
+            if call:
+                calls.append(call)
+
+        for child in node.children:
+            # Don't recurse into nested functions/classes
+            if child.type not in (
+                "function_declaration",
+                "function_expression",
+                "arrow_function",
+                "class_declaration",
+                "class_expression",
+            ):
+                self._extract_calls_recursive(child, caller_qn, calls)
+
+    def _extract_single_call(
+        self, call_node: Node, caller_qn: str
+    ) -> ExtractedCall | None:
+        """Extract a single call expression."""
+        func_node = call_node.child_by_field_name("function")
+        if not func_node:
+            return None
+
+        line = call_node.start_point[0] + 1
+        call_text = self.get_node_text(func_node)
+
+        receiver = None
+        name = call_text
+
+        if func_node.type == "member_expression":
+            # obj.method() or obj?.method()
+            obj_node = func_node.child_by_field_name("object")
+            prop_node = func_node.child_by_field_name("property")
+            if obj_node and prop_node:
+                receiver = self.get_node_text(obj_node)
+                name = self.get_node_text(prop_node)
+
+        # Extract arguments
+        arguments: list[str] = []
+        args_node = call_node.child_by_field_name("arguments")
+        if args_node:
+            for child in args_node.children:
+                if child.type not in ("(", ")", ","):
+                    arguments.append(self.get_node_text(child))
+
+        return ExtractedCall(
+            name=name,
+            line=line,
+            caller_qualified_name=caller_qn,
+            receiver=receiver,
+            arguments=arguments,
+        )
+
+    def extract_imports(self, root: Node) -> list[ExtractedImport]:
+        """Extract all import statements."""
+        imports: list[ExtractedImport] = []
+        self._extract_imports_recursive(root, imports)
+        return imports
+
+    def _extract_imports_recursive(
+        self, node: Node, imports: list[ExtractedImport]
+    ) -> None:
+        """Recursively extract import statements."""
+        if node.type == "import_statement":
+            self._extract_es6_import(node, imports)
+        elif node.type in ("lexical_declaration", "variable_declaration"):
+            # Check for require() calls
+            self._extract_require_import(node, imports)
+        else:
+            for child in node.children:
+                self._extract_imports_recursive(child, imports)
+
+    def _extract_es6_import(
+        self, node: Node, imports: list[ExtractedImport]
+    ) -> None:
+        """Extract ES6 import statement."""
+        line = node.start_point[0] + 1
+
+        # Get the source module
+        source_node = node.child_by_field_name("source")
+        if not source_node:
+            return
+
+        # Remove quotes from module path
+        module = self.get_node_text(source_node).strip("'\"")
+
+        # Check for different import patterns
+        for child in node.children:
+            if child.type == "import_clause":
+                self._extract_import_clause(child, module, line, imports)
+            elif child.type == "namespace_import":
+                # import * as name from 'module'
+                for c in child.children:
+                    if c.type == "identifier":
+                        imports.append(
+                            ExtractedImport(
+                                module=module,
+                                alias=self.get_node_text(c),
+                                is_wildcard=True,
+                                line=line,
+                            )
+                        )
+
+    def _extract_import_clause(
+        self, clause: Node, module: str, line: int, imports: list[ExtractedImport]
+    ) -> None:
+        """Extract imports from import clause."""
+        for child in clause.children:
+            if child.type == "identifier":
+                # Default import: import foo from 'module'
+                imports.append(
+                    ExtractedImport(
+                        module=module,
+                        name="default",
+                        alias=self.get_node_text(child),
+                        line=line,
+                    )
+                )
+            elif child.type == "named_imports":
+                # Named imports: import { foo, bar } from 'module'
+                for spec in child.children:
+                    if spec.type == "import_specifier":
+                        name_node = spec.child_by_field_name("name")
+                        alias_node = spec.child_by_field_name("alias")
+                        if name_node:
+                            name = self.get_node_text(name_node)
+                            alias = (
+                                self.get_node_text(alias_node) if alias_node else None
+                            )
+                            imports.append(
+                                ExtractedImport(
+                                    module=module, name=name, alias=alias, line=line
+                                )
+                            )
+            elif child.type == "namespace_import":
+                # import * as name
+                for c in child.children:
+                    if c.type == "identifier":
+                        imports.append(
+                            ExtractedImport(
+                                module=module,
+                                alias=self.get_node_text(c),
+                                is_wildcard=True,
+                                line=line,
+                            )
+                        )
+
+    def _extract_require_import(
+        self, decl_node: Node, imports: list[ExtractedImport]
+    ) -> None:
+        """Extract CommonJS require() imports."""
+        line = decl_node.start_point[0] + 1
+
+        for declarator in self._get_declarators(decl_node):
+            name_node = declarator.child_by_field_name("name")
+            value_node = declarator.child_by_field_name("value")
+
+            if not value_node or value_node.type != "call_expression":
+                continue
+
+            # Check if it's a require() call
+            func_node = value_node.child_by_field_name("function")
+            if not func_node or self.get_node_text(func_node) != "require":
+                continue
+
+            # Get the module path
+            args_node = value_node.child_by_field_name("arguments")
+            if not args_node:
+                continue
+
+            for arg in args_node.children:
+                if arg.type == "string":
+                    module = self.get_node_text(arg).strip("'\"")
+
+                    if name_node:
+                        if name_node.type == "identifier":
+                            # const foo = require('module')
+                            imports.append(
+                                ExtractedImport(
+                                    module=module,
+                                    alias=self.get_node_text(name_node),
+                                    line=line,
+                                )
+                            )
+                        elif name_node.type == "object_pattern":
+                            # const { foo, bar } = require('module')
+                            for prop in name_node.children:
+                                if prop.type == "shorthand_property_identifier_pattern":
+                                    name = self.get_node_text(prop)
+                                    imports.append(
+                                        ExtractedImport(
+                                            module=module, name=name, line=line
+                                        )
+                                    )
+                                elif prop.type == "pair_pattern":
+                                    key_node = prop.child_by_field_name("key")
+                                    val_node = prop.child_by_field_name("value")
+                                    if key_node and val_node:
+                                        imports.append(
+                                            ExtractedImport(
+                                                module=module,
+                                                name=self.get_node_text(key_node),
+                                                alias=self.get_node_text(val_node),
+                                                line=line,
+                                            )
+                                        )
+                    break