coderay 1.0.0__py3-none-any.whl

coderay/__init__.py

@@ -0,0 +1 @@
+__version__ = "1.0.0"

coderay/chunking/chunker.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+from coderay.chunking.registry import LanguageConfig, get_language_for_file
+from coderay.core.models import Chunk
+
+logger = logging.getLogger(__name__)
+
+
+def _get_symbol_name(node, source_bytes: bytes) -> str:
+    """Extract symbol name from a definition node."""
+    if node.type == "decorated_definition":
+        for child in node.children:
+            if child.type != "decorator":
+                return _get_symbol_name(child, source_bytes)
+        return ""
+
+    for child in node.children:
+        if child.type == "identifier":
+            return source_bytes[child.start_byte : child.end_byte].decode(
+                "utf-8", errors="replace"
+            )
+        if child.type in ("class", "def", "func", "function", "type"):
+            for sibling in node.children:
+                if sibling.type == "identifier":
+                    return source_bytes[sibling.start_byte : sibling.end_byte].decode(
+                        "utf-8", errors="replace"
+                    )
+    if node.type in ("property_identifier", "field_identifier"):
+        return source_bytes[node.start_byte : node.end_byte].decode(
+            "utf-8", errors="replace"
+        )
+    return ""
+
+
+def _collect_preamble_lines(
+    root, source_bytes: bytes, chunk_types: tuple[str, ...]
+) -> list[str]:
+    """Collect top-level lines that are NOT part of any chunk_type definition."""
+    lines: list[str] = []
+    for child in root.children:
+        if child.type in chunk_types:
+            continue
+        text = (
+            source_bytes[child.start_byte : child.end_byte]
+            .decode("utf-8", errors="replace")
+            .strip()
+        )
+        if text:
+            lines.append(text)
+    return lines
+
+
+def _chunk_file_with_config(
+    path: str,
+    content: str,
+    lang_cfg: LanguageConfig,
+) -> list[Chunk]:
+    """Chunk a file using the provided language configuration."""
+    try:
+        parser = lang_cfg.get_parser()
+    except Exception as e:
+        logger.warning("Could not load parser for %s (%s): %s", path, lang_cfg.name, e)
+        return []
+
+    source_bytes = content.encode("utf-8")
+    tree = parser.parse(source_bytes)
+    root = tree.root_node
+    chunks: list[Chunk] = []
+
+    def dfs(node) -> None:
+        if node.type in lang_cfg.chunk_types:
+            # [py] Avoid duplicates on decorated functions.
+            # [py] Decorators are stored with symbol of the function that is decorating
+            # [py] But the content field of the decorated function will capture them
+            if node.parent and node.parent.type in lang_cfg.chunk_types:
+                for child in node.children:
+                    dfs(child)
+                return
+            start_line = node.start_point[0] + 1
+            end_line = node.end_point[0] + 1
+            text = source_bytes[node.start_byte : node.end_byte].decode(
+                "utf-8", errors="replace"
+            )
+            symbol = _get_symbol_name(node, source_bytes) or f"<{node.type}>"
+            chunks.append(
+                Chunk(
+                    path=path,
+                    start_line=start_line,
+                    end_line=end_line,
+                    symbol=symbol,
+                    language=lang_cfg.name,
+                    content=text,
+                )
+            )
+        for child in node.children:
+            dfs(child)
+
+    if preamble_lines := _collect_preamble_lines(
+        root, source_bytes, lang_cfg.chunk_types
+    ):
+        chunks.append(
+            Chunk(
+                path=path,
+                start_line=1,
+                end_line=root.end_point[0] + 1,
+                symbol="<module>",
+                language=lang_cfg.name,
+                content="\n".join(preamble_lines),
+            ),
+        )
+
+    dfs(root)
+
+    logger.debug("Chunked %s: %d chunks", path, len(chunks))
+    return chunks
+
+
+def chunk_file(path: str | Path, content: str, language: str = "python") -> list[Chunk]:
+    """Chunk a source file into semantic units (functions, classes, preamble)."""
+    path_str = str(path) if isinstance(path, Path) else path
+    if not (lang_cfg := get_language_for_file(path_str)):
+        logger.warning("No language config for %s ", path_str)
+        return []
+    return _chunk_file_with_config(path_str, content, lang_cfg)

coderay/chunking/registry.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from tree_sitter import Language, Parser
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LanguageConfig:
+    """Configuration for a single language's tree-sitter grammar."""
+
+    name: str
+    extensions: tuple[str, ...]
+    language_fn: Callable[[], Any]
+    chunk_types: tuple[str, ...]
+    scope_types: tuple[str, ...] = ("function_definition", "class_definition")
+    import_types: tuple[str, ...] = ("import_statement", "import_from_statement")
+    call_types: tuple[str, ...] = ("call", "call_expression")
+    function_scope_types: tuple[str, ...] = ("function_definition",)
+    class_scope_types: tuple[str, ...] = ("class_definition",)
+    init_filenames: tuple[str, ...] = ()
+
+    def get_parser(self) -> Parser:
+        """Create and return a tree-sitter Parser for this language."""
+        lang = Language(self.language_fn())
+        parser = Parser(lang)
+        return parser
+
+
+def _python_language():
+    import tree_sitter_python as tspython
+
+    return tspython.language()
+
+
+def _javascript_language():
+    import tree_sitter_javascript as tsjs
+
+    return tsjs.language()
+
+
+def _typescript_language():
+    import tree_sitter_typescript as tsts
+
+    return tsts.language()
+
+
+def _go_language():
+    import tree_sitter_go as tsgo
+
+    return tsgo.language()
+
+
+PYTHON_CONFIG = LanguageConfig(
+    name="python",
+    extensions=(".py", ".pyi"),
+    language_fn=_python_language,
+    chunk_types=(
+        "function_definition",
+        "class_definition",
+        "decorated_definition",
+    ),
+    scope_types=("function_definition", "class_definition"),
+    import_types=("import_statement", "import_from_statement"),
+    call_types=("call",),
+    function_scope_types=("function_definition",),
+    class_scope_types=("class_definition",),
+    init_filenames=("__init__",),
+)
+
+JAVASCRIPT_CONFIG = LanguageConfig(
+    name="javascript",
+    extensions=(".js", ".jsx", ".mjs", ".cjs"),
+    language_fn=_javascript_language,
+    chunk_types=(
+        "function_declaration",
+        "class_declaration",
+        "method_definition",
+        "arrow_function",
+        "export_statement",
+        "lexical_declaration",
+    ),
+    scope_types=("function_declaration", "class_declaration", "method_definition"),
+    import_types=("import_statement",),
+    call_types=("call_expression",),
+    function_scope_types=("function_declaration", "method_definition"),
+    class_scope_types=("class_declaration",),
+    init_filenames=("index",),
+)
+
+TYPESCRIPT_CONFIG = LanguageConfig(
+    name="typescript",
+    extensions=(".ts", ".tsx"),
+    language_fn=_typescript_language,
+    chunk_types=(
+        "function_declaration",
+        "class_declaration",
+        "method_definition",
+        "arrow_function",
+        "export_statement",
+        "lexical_declaration",
+        "interface_declaration",
+        "type_alias_declaration",
+    ),
+    scope_types=(
+        "function_declaration",
+        "class_declaration",
+        "method_definition",
+        "interface_declaration",
+    ),
+    import_types=("import_statement",),
+    call_types=("call_expression",),
+    function_scope_types=("function_declaration", "method_definition"),
+    class_scope_types=("class_declaration", "interface_declaration"),
+    init_filenames=("index",),
+)
+
+GO_CONFIG = LanguageConfig(
+    name="go",
+    extensions=(".go",),
+    language_fn=_go_language,
+    chunk_types=(
+        "function_declaration",
+        "method_declaration",
+        "type_declaration",
+    ),
+    scope_types=("function_declaration", "method_declaration"),
+    import_types=("import_declaration",),
+    call_types=("call_expression",),
+    function_scope_types=("function_declaration", "method_declaration"),
+    class_scope_types=(),
+    init_filenames=(),
+)
+
+LANGUAGE_REGISTRY: dict[str, LanguageConfig] = {
+    "python": PYTHON_CONFIG,
+    "javascript": JAVASCRIPT_CONFIG,
+    "typescript": TYPESCRIPT_CONFIG,
+    "go": GO_CONFIG,
+}
+
+_EXTENSION_MAP: dict[str, str] = {}
+for _lang_name, _cfg in LANGUAGE_REGISTRY.items():
+    for _ext in _cfg.extensions:
+        _EXTENSION_MAP[_ext] = _lang_name
+
+
+def get_language_for_file(path: str | Path) -> LanguageConfig | None:
+    """Return the LanguageConfig for a file based on its extension, or None."""
+    ext = Path(path).suffix.lower()
+    lang_name = _EXTENSION_MAP.get(ext)
+    if lang_name is None:
+        return None
+    return LANGUAGE_REGISTRY.get(lang_name)
+
+
+def get_supported_extensions() -> set[str]:
+    """Return all file extensions we can index."""
+    return set(_EXTENSION_MAP.keys())
+
+
+def get_init_filenames() -> set[str]:
+    """Return all init-style filenames across languages (e.g. __init__, index)."""
+    names: set[str] = set()
+    for cfg in LANGUAGE_REGISTRY.values():
+        names.update(cfg.init_filenames)
+    return names
+
+
+def get_resolution_suffixes() -> list[str]:
+    """Return file suffixes for resolving import targets."""
+    suffixes: list[str] = []
+    seen: set[str] = set()
+    for cfg in LANGUAGE_REGISTRY.values():
+        for ext in cfg.extensions:
+            if ext not in seen:
+                suffixes.append(ext)
+                seen.add(ext)
+            for init in cfg.init_filenames:
+                combo = f"/{init}{ext}"
+                if combo not in seen:
+                    suffixes.append(combo)
+                    seen.add(combo)
+    return suffixes

coderay/cli/__init__.py

@@ -0,0 +1,3 @@
+from coderay.cli.commands import cli, main
+
+__all__ = ["cli", "main"]