Merge pull request #5639 from Textualize/tree-sitter-improvements

TextArea improvements - lazy language import and allow installation of only required languages

Author: willmcgugan

CHANGELOG.md

@@ -32,6 +32,10 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Tabs now accept Content or content markup https://github.com/Textualize/textual/pull/5657
 - Buttons will now use Textual markup rather than console markup
 
+### Changed
+
+- tree-sitter languages are now loaded lazily, improving cold-start time https://github.com/Textualize/textual/pull/5639
+
 ## [2.1.2] - 2025-02-26
 
 ### Fixed

src/textual/_tree_sitter.py

@@ -1,49 +1,44 @@
 from __future__ import annotations
+from importlib import import_module
+
+from textual import log
+
 
 try:
-    import tree_sitter_bash
-    import tree_sitter_css
-    import tree_sitter_go
-    import tree_sitter_html
-    import tree_sitter_java
-    import tree_sitter_javascript
-    import tree_sitter_json
-    import tree_sitter_markdown
-    import tree_sitter_python
-    import tree_sitter_regex
-    import tree_sitter_rust
-    import tree_sitter_sql
-    import tree_sitter_toml
-    import tree_sitter_xml
-    import tree_sitter_yaml
     from tree_sitter import Language
 
-    _tree_sitter = True
+    _LANGUAGE_CACHE: dict[str, Language] = {}
 
-    _languages = {
-        "python": Language(tree_sitter_python.language()),
-        "json": Language(tree_sitter_json.language()),
-        "markdown": Language(tree_sitter_markdown.language()),
-        "yaml": Language(tree_sitter_yaml.language()),
-        "toml": Language(tree_sitter_toml.language()),
-        "rust": Language(tree_sitter_rust.language()),
-        "html": Language(tree_sitter_html.language()),
-        "css": Language(tree_sitter_css.language()),
-        "xml": Language(tree_sitter_xml.language_xml()),
-        "regex": Language(tree_sitter_regex.language()),
-        "sql": Language(tree_sitter_sql.language()),
-        "javascript": Language(tree_sitter_javascript.language()),
-        "java": Language(tree_sitter_java.language()),
-        "bash": Language(tree_sitter_bash.language()),
-        "go": Language(tree_sitter_go.language()),
-    }
+    _tree_sitter = True
 
     def get_language(language_name: str) -> Language | None:
-        return _languages.get(language_name)
+        if language_name in _LANGUAGE_CACHE:
+            return _LANGUAGE_CACHE[language_name]
+
+        try:
+            module = import_module(f"tree_sitter_{language_name}")
+        except ImportError:
+            return None
+        else:
+            try:
+                if language_name == "xml":
+                    # xml uses language_xml() instead of language()
+                    # it's the only outlier amongst the languages in the `textual[syntax]` extra
+                    language = Language(module.language_xml(), name=language_name)
+                else:
+                    language = Language(module.language(), name=language_name)
+            except (OSError, AttributeError):
+                log.warning(f"Could not load language {language_name!r}.")
+                return None
+            else:
+                _LANGUAGE_CACHE[language_name] = language
+                return language
 
 except ImportError:
     _tree_sitter = False
-    _languages = {}
+
+    def get_language(language_name: str) -> Language | None:
+        return None
+
 
 TREE_SITTER = _tree_sitter
-BUILTIN_LANGUAGES: dict[str, "Language"] = _languages

src/textual/document/_syntax_aware_document.py

@@ -8,7 +8,6 @@
     TREE_SITTER = False
 
 
-from textual._tree_sitter import BUILTIN_LANGUAGES
 from textual.document._document import Document, EditResult, Location, _utf8_encode
 
 
@@ -17,59 +16,30 @@ class SyntaxAwareDocumentError(Exception):
 
 
 class SyntaxAwareDocument(Document):
-    """A wrapper around a Document which also maintains a tree-sitter syntax
+    """A subclass of Document which also maintains a tree-sitter syntax
     tree when the document is edited.
-
-    The primary reason for this split is actually to keep tree-sitter stuff separate,
-    since it isn't supported in Python 3.7. By having the tree-sitter code
-    isolated in this subclass, it makes it easier to conditionally import. However,
-    it does come with other design flaws (e.g. Document is required to have methods
-    which only really make sense on SyntaxAwareDocument).
-
-    If you're reading this and Python 3.7 is no longer supported by Textual,
-    consider merging this subclass into the `Document` superclass.
     """
 
     def __init__(
         self,
         text: str,
-        language: str | Language,
+        language: Language,
     ):
         """Construct a SyntaxAwareDocument.
 
         Args:
             text: The initial text contained in the document.
-            language: The language to use. You can pass a string to use a supported
-                language, or pass in your own tree-sitter `Language` object.
+            language: The tree-sitter language to use.
         """
 
         if not TREE_SITTER:
-            raise RuntimeError("SyntaxAwareDocument unavailable.")
+            raise RuntimeError(
+                "SyntaxAwareDocument unavailable - tree-sitter is not installed."
+            )
 
         super().__init__(text)
-        self.language: Language | None = None
-        """The tree-sitter Language or None if tree-sitter is unavailable."""
-
-        self._parser: Parser | None = None
-
-        from textual._tree_sitter import get_language
-
-        # If the language is `None`, then avoid doing any parsing related stuff.
-        if isinstance(language, str):
-            if language not in BUILTIN_LANGUAGES:
-                raise SyntaxAwareDocumentError(f"Invalid language {language!r}")
-            # If tree-sitter-languages is not installed properly, get_language
-            # and get_parser may raise an OSError when unable to load their
-            # resources
-
-            try:
-                self.language = get_language(language)
-            except OSError as e:
-                raise SyntaxAwareDocumentError(
-                    f"Could not find binaries for {language!r}"
-                ) from e
-        else:
-            self.language = language
+        self.language: Language = language
+        """The tree-sitter Language."""
 
         self._parser = Parser(self.language)
         """The tree-sitter Parser or None if tree-sitter is unavailable."""
@@ -90,16 +60,6 @@ def prepare_query(self, query: str) -> Query | None:
         Returns:
             The prepared query.
         """
-        if not TREE_SITTER:
-            raise SyntaxAwareDocumentError(
-                "Couldn't prepare query - tree-sitter is not available on this architecture."
-            )
-
-        if self.language is None:
-            raise SyntaxAwareDocumentError(
-                "Couldn't prepare query - no language assigned."
-            )
-
         return self.language.query(query)
 
     def query_syntax_tree(
@@ -122,12 +82,6 @@ def query_syntax_tree(
         Returns:
             A tuple containing the nodes and text captured by the query.
         """
-
-        if not TREE_SITTER:
-            raise SyntaxAwareDocumentError(
-                "tree-sitter is not available on this architecture."
-            )
-
         captures_kwargs = {}
         if start_point is not None:
             captures_kwargs["start_point"] = start_point