Make language frontends optional

src/docforge/treesitter_logic.py

@@ -9,18 +9,20 @@ data; it is never imported, compiled, or executed.
 from __future__ import annotations
 
 import hashlib
+import importlib
 from collections.abc import Iterable
 from dataclasses import dataclass
 from functools import lru_cache
-
-import tree_sitter_cpp
-import tree_sitter_javascript
-from tree_sitter import Language, Node, Parser
+from typing import TYPE_CHECKING
 
 from .errors import DocForgeError
 from .models import LogicEdge, LogicNode, LogicProjection
 
+if TYPE_CHECKING:
+    from tree_sitter import Language, Node
+
 _COMMENT_NODE_TYPES = frozenset({"comment"})
+_SCRIPT_LANGUAGES = frozenset({"javascript", "typescript"})
 
 
 @dataclass(frozen=True)
@@ -79,7 +81,34 @@ class _LanguageProfile:
 def _javascript_profile() -> _LanguageProfile:
     return _LanguageProfile(
         name="javascript",
-        language=Language(tree_sitter_javascript.language()),
+        language=_optional_language(
+            grammar_module="tree_sitter_javascript",
+            grammar_function="language",
+            extra="javascript",
+        ),
+        root_type="program",
+        block_types=frozenset({"program", "statement_block"}),
+        function_types=frozenset(
+            {"function_declaration", "generator_function_declaration", "method_definition"}
+        ),
+        loop_types=frozenset(
+            {"while_statement", "do_statement", "for_statement", "for_in_statement"}
+        ),
+        return_types=frozenset({"return_statement"}),
+        raise_types=frozenset({"throw_statement"}),
+        switch_case_types=frozenset({"switch_case", "switch_default"}),
+    )
+
+
+@lru_cache(maxsize=1)
+def _typescript_profile() -> _LanguageProfile:
+    return _LanguageProfile(
+        name="typescript",
+        language=_optional_language(
+            grammar_module="tree_sitter_typescript",
+            grammar_function="language_typescript",
+            extra="typescript",
+        ),
         root_type="program",
         block_types=frozenset({"program", "statement_block"}),
         function_types=frozenset(
@@ -98,7 +127,11 @@ def _javascript_profile() -> _LanguageProfile:
 def _cpp_profile() -> _LanguageProfile:
     return _LanguageProfile(
         name="cpp",
-        language=Language(tree_sitter_cpp.language()),
+        language=_optional_language(
+            grammar_module="tree_sitter_cpp",
+            grammar_function="language",
+            extra="cpp",
+        ),
         root_type="translation_unit",
         block_types=frozenset({"translation_unit", "compound_statement"}),
         function_types=frozenset({"function_definition"}),
@@ -116,6 +149,28 @@ def _cpp_profile() -> _LanguageProfile:
     )
 
 
+def _optional_language(
+    *,
+    grammar_module: str,
+    grammar_function: str,
+    extra: str,
+) -> Language:
+    try:
+        tree_sitter = importlib.import_module("tree_sitter")
+        grammar = importlib.import_module(grammar_module)
+    except ModuleNotFoundError as error:
+        raise DocForgeError(
+            "optional_dependency_missing",
+            "The requested language frontend is not installed",
+            extra=extra,
+            install=f"docforge[{extra}]",
+            missing_module=error.name,
+        ) from error
+    language_type = tree_sitter.Language
+    language_factory = getattr(grammar, grammar_function)
+    return language_type(language_factory())
+
+
 def analyze_javascript_source(
     source: str,
     *,
@@ -136,12 +191,38 @@ def analyze_javascript_source(
     )
 
 
+def analyze_typescript_source(
+    source: str,
+    *,
+    source_id: str,
+    owners: Iterable[TreeSitterLogicOwner],
+    filename: str = "<typescript-source>",
+    max_nodes_per_function: int = 2_000,
+) -> tuple[LogicProjection, ...]:
+    """Build control-flow projections for named TypeScript functions and methods."""
+
+    return _analyze_tree_sitter_source(
+        source,
+        source_id=source_id,
+        owners=owners,
+        filename=filename,
+        profile=_typescript_profile(),
+        max_nodes_per_function=max_nodes_per_function,
+    )
+
+
 def discover_javascript_functions(source: str) -> tuple[DiscoveredFunction, ...]:
     """Return named JavaScript functions, methods, and assigned arrow functions."""
 
     return _discover_functions(source, _javascript_profile())
 
 
+def discover_typescript_functions(source: str) -> tuple[DiscoveredFunction, ...]:
+    """Return named TypeScript functions, methods, and assigned arrow functions."""
+
+    return _discover_functions(source, _typescript_profile())
+
+
 def discover_cpp_functions(source: str) -> tuple[DiscoveredFunction, ...]:
     """Return named C++ functions and methods."""
 
@@ -172,6 +253,8 @@ def _discover_functions(
     source: str,
     profile: _LanguageProfile,
 ) -> tuple[DiscoveredFunction, ...]:
+    from tree_sitter import Parser
+
     raw = source.encode("utf-8")
     parser = Parser(profile.language)
     tree = parser.parse(raw)
@@ -212,6 +295,8 @@ def _analyze_tree_sitter_source(
     profile: _LanguageProfile,
     max_nodes_per_function: int,
 ) -> tuple[LogicProjection, ...]:
+    from tree_sitter import Parser
+
     if max_nodes_per_function < 2:
         raise ValueError("max_nodes_per_function must allow entry and exit nodes")
     raw = source.encode("utf-8")
@@ -293,7 +378,7 @@ def _function_definitions(
 
 
 def _scope_name(node: Node, raw: bytes, profile: _LanguageProfile) -> str | None:
-    if profile.name == "javascript" and node.type in {"class_declaration", "class"}:
+    if profile.name in _SCRIPT_LANGUAGES and node.type in {"class_declaration", "class"}:
         return _field_text(node, "name", raw)
     if profile.name == "cpp" and node.type in {
         "namespace_definition",
@@ -307,11 +392,11 @@ def _scope_name(node: Node, raw: bytes, profile: _LanguageProfile) -> str | None
 
 def _function_name(node: Node, raw: bytes, profile: _LanguageProfile) -> str | None:
     if node.type in profile.function_types:
-        if profile.name == "javascript":
+        if profile.name in _SCRIPT_LANGUAGES:
             return _field_text(node, "name", raw)
         declarator = node.child_by_field_name("declarator")
         return _declarator_name(declarator, raw) if declarator is not None else None
-    if profile.name != "javascript" or node.type != "variable_declarator":
+    if profile.name not in _SCRIPT_LANGUAGES or node.type != "variable_declarator":
         return None
     value = node.child_by_field_name("value")
     if value is None or value.type not in {"arrow_function", "function_expression"}: