codegraph-ai 0.1.0__py3-none-any.whl

codegraph/adapters/js_adapter.py

@@ -0,0 +1,556 @@
+"""JavaScript / TypeScript source code adapter using tree-sitter.
+
+Handles ``.js``, ``.mjs``, ``.cjs``, ``.ts``, ``.tsx`` and ``.jsx`` files.
+Extracts:
+  - ``function`` declarations
+  - Arrow functions assigned to ``const`` / ``let`` / ``var``
+  - Class declarations and their methods
+  - Function calls
+  - ES module imports (``import … from``) and CommonJS ``require()``
+"""
+
+from __future__ import annotations
+
+from tree_sitter_language_pack import get_parser
+
+from codegraph.adapters.base import BaseAdapter
+from codegraph.models import (
+    CallInfo,
+    ParsedClass,
+    ParsedFunction,
+    ParsedImport,
+    ParseResult,
+)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _node_text(node) -> str:
+    """Return the UTF-8 text of a tree-sitter node."""
+    return node.text.decode("utf-8") if node and node.text else ""
+
+
+def _extract_jsdoc(node) -> str:
+    """Extract a JSDoc comment preceding a node.
+
+    Looks at the previous sibling; if it is a ``comment`` whose text
+    starts with ``/**``, we strip the delimiters and leading ``*``'s.
+    """
+    prev = node.prev_sibling
+    if prev is None or prev.type != "comment":
+        # For method_definition inside class_body the comment is a sibling too
+        return ""
+    text = _node_text(prev)
+    if not text.startswith("/**"):
+        return ""
+    # Strip /** ... */
+    text = text[3:]
+    if text.endswith("*/"):
+        text = text[:-2]
+    lines = text.splitlines()
+    cleaned: list[str] = []
+    for line in lines:
+        line = line.strip()
+        if line.startswith("*"):
+            line = line[1:].strip()
+        if line.startswith("@"):
+            break  # stop at first tag
+        if line:
+            cleaned.append(line)
+    return " ".join(cleaned)
+
+
+def _build_js_signature(name: str, params_node, ret_node=None) -> str:
+    """Build a human-readable function signature."""
+    params = _node_text(params_node) if params_node else "()"
+    ret = ""
+    if ret_node:
+        ret = f": {_node_text(ret_node)}"
+    return f"function {name}{params}{ret}"
+
+
+def _collect_calls(node, calls: list[CallInfo]) -> None:
+    """Recursively collect function calls with receiver context."""
+    if node.type == "call_expression":
+        func = node.child_by_field_name("function")
+        if func:
+            if func.type == "member_expression":
+                obj_node = func.child_by_field_name("object")
+                prop_node = func.child_by_field_name("property")
+                receiver = _node_text(obj_node) if obj_node else None
+                callee = _node_text(prop_node) if prop_node else _node_text(func)
+                if receiver and "." in receiver:
+                    receiver = receiver.rsplit(".", 1)[-1]
+                if callee != "require":
+                    calls.append(CallInfo(
+                        callee_name=callee,
+                        receiver=receiver,
+                        raw_expression=_node_text(func),
+                    ))
+            else:
+                callee = _node_text(func)
+                if callee != "require":
+                    calls.append(CallInfo(
+                        callee_name=callee,
+                        receiver=None,
+                        raw_expression=callee,
+                    ))
+    for child in node.children:
+        _collect_calls(child, calls)
+
+
+# ---------------------------------------------------------------------------
+# Adapter
+# ---------------------------------------------------------------------------
+
+
+class JsAdapter(BaseAdapter):
+    """Extract functions, classes, calls and imports from JS / TS files."""
+
+    _LANG_MAP = {
+        ".js": "javascript",
+        ".mjs": "javascript",
+        ".cjs": "javascript",
+        ".jsx": "javascript",
+        ".ts": "typescript",
+        ".tsx": "tsx",
+    }
+
+    def __init__(self) -> None:
+        self._parsers: dict[str, object] = {}
+
+    def _get_parser(self, ext: str):
+        lang = self._LANG_MAP.get(ext, "typescript")
+        if lang not in self._parsers:
+            self._parsers[lang] = get_parser(lang)
+        return self._parsers[lang]
+
+    # -- BaseAdapter interface ------------------------------------------------
+
+    def language_name(self) -> str:
+        return "javascript"
+
+    def supported_extensions(self) -> list[str]:
+        return list(self._LANG_MAP.keys())
+
+    def parse_file(self, source: bytes, file_path: str) -> ParseResult:
+        ext = ""
+        for e in self._LANG_MAP:
+            if file_path.endswith(e):
+                ext = e
+                break
+        parser = self._get_parser(ext)
+        tree = parser.parse(source)
+        root = tree.root_node
+
+        functions: list[ParsedFunction] = []
+        classes: list[ParsedClass] = []
+        imports: list[ParsedImport] = []
+
+        self._walk(root, file_path, functions, classes, imports)
+        return ParseResult(functions=functions, classes=classes, imports=imports)
+
+    # -- Internal helpers -----------------------------------------------------
+
+    def _walk(
+        self,
+        node,
+        file_path: str,
+        functions: list[ParsedFunction],
+        classes: list[ParsedClass],
+        imports: list[ParsedImport],
+    ) -> None:
+        """Walk top-level children and populate result lists."""
+        for child in node.children:
+            self._process_node(
+                child, file_path, functions, classes, imports, class_name=None
+            )
+
+    def _process_node(
+        self,
+        child,
+        file_path: str,
+        functions: list[ParsedFunction],
+        classes: list[ParsedClass],
+        imports: list[ParsedImport],
+        class_name: str | None,
+    ) -> None:
+        """Handle a single AST node."""
+        # --- function declaration ---
+        if child.type == "function_declaration":
+            self._extract_function_decl(child, file_path, functions, class_name)
+
+        # --- const/let/var = arrow / function expression ---
+        elif child.type == "lexical_declaration":
+            self._extract_lexical(child, file_path, functions, imports)
+        elif child.type == "variable_declaration":
+            self._extract_lexical(child, file_path, functions, imports)
+
+        # --- class declaration ---
+        elif child.type == "class_declaration":
+            self._extract_class(child, file_path, functions, classes)
+
+        # --- export statement: unwrap and recurse ---
+        elif child.type == "export_statement":
+            for ec in child.children:
+                self._process_node(
+                    ec, file_path, functions, classes, imports, class_name
+                )
+
+        # --- imports ---
+        elif child.type == "import_statement":
+            self._extract_import(child, file_path, imports)
+
+        # --- expression_statement: may contain an exported assignment ---
+        elif child.type == "expression_statement":
+            for ec in child.children:
+                if ec.type == "assignment_expression":
+                    self._extract_assignment_func(
+                        ec, file_path, functions
+                    )
+
+    # -- Extractors -----------------------------------------------------------
+
+    def _extract_function_decl(
+        self,
+        node,
+        file_path: str,
+        functions: list[ParsedFunction],
+        class_name: str | None,
+    ) -> None:
+        """Extract a ``function foo(…) { … }`` declaration."""
+        name_node = node.child_by_field_name("name")
+        name = _node_text(name_node) if name_node else "anonymous"
+        params = node.child_by_field_name("parameters")
+        ret = node.child_by_field_name("return_type")
+        body = node.child_by_field_name("body")
+
+        sig = _build_js_signature(name, params, ret)
+        doc = _extract_jsdoc(node)
+        calls: list[CallInfo] = []
+        if body:
+            _collect_calls(body, calls)
+
+        start_line = node.start_point[0] + 1
+        end_line = node.end_point[0] + 1
+        qualified = (
+            f"{file_path}:{class_name}.{name}" if class_name
+            else f"{file_path}:{name}"
+        )
+
+        functions.append(
+            ParsedFunction(
+                name=name,
+                qualified_name=qualified,
+                signature=sig,
+                file_path=file_path,
+                start_line=start_line,
+                end_line=end_line,
+                doc_comment=doc,
+                call_names=[c.callee_name for c in calls],
+                calls=calls,
+                class_name=class_name,
+            )
+        )
+
+    def _extract_lexical(
+        self,
+        node,
+        file_path: str,
+        functions: list[ParsedFunction],
+        imports: list[ParsedImport],
+    ) -> None:
+        """Handle ``const x = …`` — may be arrow function or require()."""
+        for child in node.children:
+            if child.type != "variable_declarator":
+                continue
+            name_node = child.child_by_field_name("name")
+            value_node = child.child_by_field_name("value")
+            if name_node is None or value_node is None:
+                continue
+
+            name = _node_text(name_node)
+
+            # Arrow function: const foo = (…) => { … }
+            if value_node.type == "arrow_function":
+                self._extract_arrow(
+                    name, value_node, node, file_path, functions
+                )
+            # Function expression: const foo = function(…) { … }
+            elif value_node.type == "function_expression":
+                self._extract_func_expr(
+                    name, value_node, node, file_path, functions
+                )
+            # require(): const x = require("mod")
+            elif value_node.type == "call_expression":
+                func_node = value_node.child_by_field_name("function")
+                if func_node and _node_text(func_node) == "require":
+                    args = value_node.child_by_field_name("arguments")
+                    if args:
+                        for a in args.children:
+                            if a.type == "string":
+                                mod = _node_text(a).strip("'\"")
+                                imports.append(
+                                    ParsedImport(
+                                        source_path=file_path,
+                                        target_module=mod,
+                                    )
+                                )
+            # await import() — dynamic import
+            elif value_node.type == "await_expression":
+                pass  # skip dynamic imports for now
+
+    def _extract_arrow(
+        self,
+        name: str,
+        arrow_node,
+        decl_node,
+        file_path: str,
+        functions: list[ParsedFunction],
+    ) -> None:
+        """Extract an arrow function assigned to a variable."""
+        params = arrow_node.child_by_field_name("parameters")
+        ret = arrow_node.child_by_field_name("return_type")
+        body = arrow_node.child_by_field_name("body")
+
+        sig = _build_js_signature(name, params, ret)
+        doc = _extract_jsdoc(decl_node)
+        calls: list[CallInfo] = []
+        if body:
+            _collect_calls(body, calls)
+
+        start_line = decl_node.start_point[0] + 1
+        end_line = decl_node.end_point[0] + 1
+
+        functions.append(
+            ParsedFunction(
+                name=name,
+                qualified_name=f"{file_path}:{name}",
+                signature=sig,
+                file_path=file_path,
+                start_line=start_line,
+                end_line=end_line,
+                doc_comment=doc,
+                call_names=[c.callee_name for c in calls],
+                calls=calls,
+                class_name=None,
+            )
+        )
+
+    def _extract_func_expr(
+        self,
+        name: str,
+        func_node,
+        decl_node,
+        file_path: str,
+        functions: list[ParsedFunction],
+    ) -> None:
+        """Extract a function expression assigned to a variable."""
+        params = func_node.child_by_field_name("parameters")
+        ret = func_node.child_by_field_name("return_type")
+        body = func_node.child_by_field_name("body")
+
+        sig = _build_js_signature(name, params, ret)
+        doc = _extract_jsdoc(decl_node)
+        calls: list[CallInfo] = []
+        if body:
+            _collect_calls(body, calls)
+
+        start_line = decl_node.start_point[0] + 1
+        end_line = decl_node.end_point[0] + 1
+
+        functions.append(
+            ParsedFunction(
+                name=name,
+                qualified_name=f"{file_path}:{name}",
+                signature=sig,
+                file_path=file_path,
+                start_line=start_line,
+                end_line=end_line,
+                doc_comment=doc,
+                call_names=[c.callee_name for c in calls],
+                calls=calls,
+                class_name=None,
+            )
+        )
+
+    def _extract_class(
+        self,
+        node,
+        file_path: str,
+        functions: list[ParsedFunction],
+        classes: list[ParsedClass],
+    ) -> None:
+        """Extract a class and its methods."""
+        name_node = node.child_by_field_name("name")
+        cls_name = _node_text(name_node) if name_node else "AnonymousClass"
+        start_line = node.start_point[0] + 1
+        end_line = node.end_point[0] + 1
+        qualified = f"{file_path}:{cls_name}"
+
+        base_classes: list[str] = []
+        for child in node.children:
+            if child.type == "class_heritage":
+                for hc in child.children:
+                    if hc.is_named:
+                        base_classes.append(_node_text(hc))
+                        break
+
+        method_names: list[str] = []
+        body = node.child_by_field_name("body")
+        if body:
+            for child in body.children:
+                if child.type == "method_definition":
+                    m_name = self._method_name(child)
+                    method_names.append(m_name)
+                    self._extract_method(
+                        child, file_path, functions, cls_name
+                    )
+                elif child.type == "public_field_definition":
+                    # class field with arrow function initializer
+                    fname = child.child_by_field_name("name")
+                    val = child.child_by_field_name("value")
+                    if fname and val and val.type == "arrow_function":
+                        m_name = _node_text(fname)
+                        method_names.append(m_name)
+                        self._extract_arrow(
+                            m_name, val, child, file_path, functions
+                        )
+                        # Fix class_name on the last appended function
+                        functions[-1] = ParsedFunction(
+                            **{
+                                **functions[-1].__dict__,
+                                "class_name": cls_name,
+                                "qualified_name": f"{file_path}:{cls_name}.{m_name}",
+                            }
+                        )
+
+        classes.append(
+            ParsedClass(
+                name=cls_name,
+                qualified_name=qualified,
+                file_path=file_path,
+                start_line=start_line,
+                end_line=end_line,
+                method_names=method_names,
+                base_classes=base_classes,
+            )
+        )
+
+    def _extract_method(
+        self,
+        node,
+        file_path: str,
+        functions: list[ParsedFunction],
+        class_name: str,
+    ) -> None:
+        """Extract a method_definition inside a class body."""
+        name = self._method_name(node)
+        params = node.child_by_field_name("parameters")
+        ret = node.child_by_field_name("return_type")
+        body = node.child_by_field_name("body")
+
+        sig = _build_js_signature(name, params, ret)
+        doc = _extract_jsdoc(node)
+        calls: list[CallInfo] = []
+        if body:
+            _collect_calls(body, calls)
+
+        start_line = node.start_point[0] + 1
+        end_line = node.end_point[0] + 1
+
+        functions.append(
+            ParsedFunction(
+                name=name,
+                qualified_name=f"{file_path}:{class_name}.{name}",
+                signature=sig,
+                file_path=file_path,
+                start_line=start_line,
+                end_line=end_line,
+                doc_comment=doc,
+                call_names=[c.callee_name for c in calls],
+                calls=calls,
+                class_name=class_name,
+            )
+        )
+
+    def _extract_import(
+        self,
+        node,
+        file_path: str,
+        imports: list[ParsedImport],
+    ) -> None:
+        """Extract ES module import with imported names and relative detection."""
+        source_node = node.child_by_field_name("source")
+        if not source_node:
+            return
+
+        mod = _node_text(source_node).strip("'\"")
+        is_relative = mod.startswith(".")
+
+        imported_names: list[str] = []
+        for child in node.children:
+            if child.type == "import_clause":
+                self._extract_import_clause_names(child, imported_names)
+
+        imports.append(
+            ParsedImport(
+                source_path=file_path,
+                target_module=mod,
+                imported_names=imported_names,
+                is_relative=is_relative,
+            )
+        )
+
+    @staticmethod
+    def _extract_import_clause_names(
+        clause_node, names: list[str]
+    ) -> None:
+        """Extract bound names from an import_clause node."""
+        for child in clause_node.children:
+            if child.type == "identifier":
+                names.append(_node_text(child))
+            elif child.type == "named_imports":
+                for spec in child.children:
+                    if spec.type == "import_specifier":
+                        name_node = spec.child_by_field_name("name")
+                        if name_node:
+                            names.append(_node_text(name_node))
+            elif child.type == "namespace_import":
+                for nc in child.children:
+                    if nc.type == "identifier":
+                        names.append(_node_text(nc))
+                        break
+
+    def _extract_assignment_func(
+        self,
+        node,
+        file_path: str,
+        functions: list[ParsedFunction],
+    ) -> None:
+        """Handle ``module.exports = function …`` or bare assignment arrows."""
+        right = node.child_by_field_name("right")
+        left = node.child_by_field_name("left")
+        if right is None or left is None:
+            return
+
+        name = _node_text(left)
+        # Simplify: module.exports -> exports
+        if "." in name:
+            name = name.rsplit(".", 1)[-1]
+
+        if right.type == "arrow_function":
+            self._extract_arrow(name, right, node, file_path, functions)
+        elif right.type == "function_expression":
+            self._extract_func_expr(name, right, node, file_path, functions)
+
+    # -- Tiny helpers ---------------------------------------------------------
+
+    @staticmethod
+    def _method_name(method_node) -> str:
+        """Return the name of a method_definition node."""
+        name_node = method_node.child_by_field_name("name")
+        return _node_text(name_node) if name_node else "anonymous"