administrator/DocForge2
1
0
Fork 0
Code Issues Pull requests Projects Releases 2 Packages Wiki Activity Actions Pages

Make language frontends optional

Browse source
This commit is contained in:
Andraxion 2026-07-29 14:18:31 -04:00
parent fee2bb0080
commit 395d732348
4 changed files with 207 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

21
pyproject.toml
View file

@ -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]

103
src/docforge/treesitter_logic.py
View file

@ -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"}:

41
tests/test_treesitter_logic.py
View file

@ -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()

54
uv.lock generated
View file

@ -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"