diff --git a/pyproject.toml b/pyproject.toml index 6024c3e..e3c16be 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -13,9 +13,26 @@ authors = [{ name = "Worldforge contributors" }] dependencies = [ "markdown-it-py>=4.2,<5", "mcp>=1.28,<2", +] + +[project.optional-dependencies] +javascript = [ + "tree-sitter>=0.25,<0.26", + "tree-sitter-javascript>=0.25,<0.26", +] +typescript = [ + "tree-sitter>=0.25,<0.26", + "tree-sitter-typescript>=0.23,<0.24", +] +cpp = [ + "tree-sitter>=0.25,<0.26", + "tree-sitter-cpp>=0.23,<0.24", +] +languages = [ "tree-sitter>=0.25,<0.26", "tree-sitter-cpp>=0.23,<0.24", "tree-sitter-javascript>=0.25,<0.26", + "tree-sitter-typescript>=0.23,<0.24", ] [dependency-groups] @@ -23,6 +40,10 @@ dev = [ "jsonschema>=4.25,<5", "pytest>=9.1,<10", "ruff>=0.15,<1", + "tree-sitter>=0.25,<0.26", + "tree-sitter-cpp>=0.23,<0.24", + "tree-sitter-javascript>=0.25,<0.26", + "tree-sitter-typescript>=0.23,<0.24", ] [project.scripts] diff --git a/src/docforge/treesitter_logic.py b/src/docforge/treesitter_logic.py index 41f981a..c6892bb 100644 --- a/src/docforge/treesitter_logic.py +++ b/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 = "", + 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"}: diff --git a/tests/test_treesitter_logic.py b/tests/test_treesitter_logic.py index 51913b3..dad4da3 100644 --- a/tests/test_treesitter_logic.py +++ b/tests/test_treesitter_logic.py @@ -7,6 +7,8 @@ from docforge.treesitter_logic import ( TreeSitterLogicOwner, analyze_cpp_source, analyze_javascript_source, + analyze_typescript_source, + discover_typescript_functions, ) @@ -200,5 +202,44 @@ public: self.assertEqual(missing.exception.code, "missing_logic_owner") +class TypeScriptLogicTests(unittest.TestCase): + def test_typed_function_and_method_use_the_typescript_grammar(self) -> None: + source = """ +interface Choice { + enabled: boolean; +} + +function choose(choice: Choice): number { + if (choice.enabled) { + return 1; + } + return 0; +} + +class Worker { + run(choice: Choice): number { + return choose(choice); + } +} +""".strip() + + discovered = discover_typescript_functions(source) + self.assertEqual( + [("Worker.run", 13), ("choose", 5)], + [(item.qualified_name, item.line) for item in discovered], + ) + projection = analyze_typescript_source( + source, + source_id="source.typescript", + owners=(TreeSitterLogicOwner("ts.symbol.choose", "choose", 5),), + )[0] + + self.assertIn("choice.enabled", {node.label for node in projection.nodes}) + self.assertEqual( + {"return"}, + {edge.relation for edge in projection.edges if edge.relation == "return"}, + ) + + if __name__ == "__main__": unittest.main() diff --git a/uv.lock b/uv.lock index 21b0bfd..de8ab13 100644 --- a/uv.lock +++ b/uv.lock @@ -211,9 +211,26 @@ source = { editable = "." } dependencies = [ { name = "markdown-it-py" }, { name = "mcp" }, +] + +[package.optional-dependencies] +cpp = [ + { name = "tree-sitter" }, + { name = "tree-sitter-cpp" }, +] +javascript = [ + { name = "tree-sitter" }, + { name = "tree-sitter-javascript" }, +] +languages = [ { name = "tree-sitter" }, { name = "tree-sitter-cpp" }, { name = "tree-sitter-javascript" }, + { name = "tree-sitter-typescript" }, +] +typescript = [ + { name = "tree-sitter" }, + { name = "tree-sitter-typescript" }, ] [package.dev-dependencies] @@ -221,22 +238,38 @@ dev = [ { name = "jsonschema" }, { name = "pytest" }, { name = "ruff" }, + { name = "tree-sitter" }, + { name = "tree-sitter-cpp" }, + { name = "tree-sitter-javascript" }, + { name = "tree-sitter-typescript" }, ] [package.metadata] requires-dist = [ { name = "markdown-it-py", specifier = ">=4.2,<5" }, { name = "mcp", specifier = ">=1.28,<2" }, - { name = "tree-sitter", specifier = ">=0.25,<0.26" }, - { name = "tree-sitter-cpp", specifier = ">=0.23,<0.24" }, - { name = "tree-sitter-javascript", specifier = ">=0.25,<0.26" }, + { name = "tree-sitter", marker = "extra == 'cpp'", specifier = ">=0.25,<0.26" }, + { name = "tree-sitter", marker = "extra == 'javascript'", specifier = ">=0.25,<0.26" }, + { name = "tree-sitter", marker = "extra == 'languages'", specifier = ">=0.25,<0.26" }, + { name = "tree-sitter", marker = "extra == 'typescript'", specifier = ">=0.25,<0.26" }, + { name = "tree-sitter-cpp", marker = "extra == 'cpp'", specifier = ">=0.23,<0.24" }, + { name = "tree-sitter-cpp", marker = "extra == 'languages'", specifier = ">=0.23,<0.24" }, + { name = "tree-sitter-javascript", marker = "extra == 'javascript'", specifier = ">=0.25,<0.26" }, + { name = "tree-sitter-javascript", marker = "extra == 'languages'", specifier = ">=0.25,<0.26" }, + { name = "tree-sitter-typescript", marker = "extra == 'languages'", specifier = ">=0.23,<0.24" }, + { name = "tree-sitter-typescript", marker = "extra == 'typescript'", specifier = ">=0.23,<0.24" }, ] +provides-extras = ["javascript", "typescript", "cpp", "languages"] [package.metadata.requires-dev] dev = [ { name = "jsonschema", specifier = ">=4.25,<5" }, { name = "pytest", specifier = ">=9.1,<10" }, { name = "ruff", specifier = ">=0.15,<1" }, + { name = "tree-sitter", specifier = ">=0.25,<0.26" }, + { name = "tree-sitter-cpp", specifier = ">=0.23,<0.24" }, + { name = "tree-sitter-javascript", specifier = ">=0.25,<0.26" }, + { name = "tree-sitter-typescript", specifier = ">=0.23,<0.24" }, ] [[package]] @@ -804,6 +837,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/2e/1f/f9eba1038b7d4394410f3c0a6ec2122b590cd7acb03f196e52fa57ebbe72/tree_sitter_javascript-0.25.0-cp310-abi3-win_arm64.whl", hash = "sha256:622a69d677aa7f6ee2931d8c77c981a33f0ebb6d275aa9d43d3397c879a9bb0b", size = 61668, upload-time = "2025-09-01T07:13:43.803Z" }, ] +[[package]] +name = "tree-sitter-typescript" +version = "0.23.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1e/fc/bb52958f7e399250aee093751e9373a6311cadbe76b6e0d109b853757f35/tree_sitter_typescript-0.23.2.tar.gz", hash = "sha256:7b167b5827c882261cb7a50dfa0fb567975f9b315e87ed87ad0a0a3aedb3834d", size = 773053, upload-time = "2024-11-11T02:36:11.396Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/28/95/4c00680866280e008e81dd621fd4d3f54aa3dad1b76b857a19da1b2cc426/tree_sitter_typescript-0.23.2-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:3cd752d70d8e5371fdac6a9a4df9d8924b63b6998d268586f7d374c9fba2a478", size = 286677, upload-time = "2024-11-11T02:35:58.839Z" }, + { url = "https://files.pythonhosted.org/packages/8f/2f/1f36fda564518d84593f2740d5905ac127d590baf5c5753cef2a88a89c15/tree_sitter_typescript-0.23.2-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:c7cc1b0ff5d91bac863b0e38b1578d5505e718156c9db577c8baea2557f66de8", size = 302008, upload-time = "2024-11-11T02:36:00.733Z" }, + { url = "https://files.pythonhosted.org/packages/96/2d/975c2dad292aa9994f982eb0b69cc6fda0223e4b6c4ea714550477d8ec3a/tree_sitter_typescript-0.23.2-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4b1eed5b0b3a8134e86126b00b743d667ec27c63fc9de1b7bb23168803879e31", size = 351987, upload-time = "2024-11-11T02:36:02.669Z" }, + { url = "https://files.pythonhosted.org/packages/49/d1/a71c36da6e2b8a4ed5e2970819b86ef13ba77ac40d9e333cb17df6a2c5db/tree_sitter_typescript-0.23.2-cp39-abi3-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e96d36b85bcacdeb8ff5c2618d75593ef12ebaf1b4eace3477e2bdb2abb1752c", size = 344960, upload-time = "2024-11-11T02:36:04.443Z" }, + { url = "https://files.pythonhosted.org/packages/7f/cb/f57b149d7beed1a85b8266d0c60ebe4c46e79c9ba56bc17b898e17daf88e/tree_sitter_typescript-0.23.2-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8d4f0f9bcb61ad7b7509d49a1565ff2cc363863644a234e1e0fe10960e55aea0", size = 340245, upload-time = "2024-11-11T02:36:06.473Z" }, + { url = "https://files.pythonhosted.org/packages/8b/ab/dd84f0e2337296a5f09749f7b5483215d75c8fa9e33738522e5ed81f7254/tree_sitter_typescript-0.23.2-cp39-abi3-win_amd64.whl", hash = "sha256:3f730b66396bc3e11811e4465c41ee45d9e9edd6de355a58bbbc49fa770da8f9", size = 278015, upload-time = "2024-11-11T02:36:07.631Z" }, + { url = "https://files.pythonhosted.org/packages/9f/e4/81f9a935789233cf412a0ed5fe04c883841d2c8fb0b7e075958a35c65032/tree_sitter_typescript-0.23.2-cp39-abi3-win_arm64.whl", hash = "sha256:05db58f70b95ef0ea126db5560f3775692f609589ed6f8dd0af84b7f19f1cbb7", size = 274052, upload-time = "2024-11-11T02:36:09.514Z" }, +] + [[package]] name = "typing-extensions" version = "4.16.0"