Fix issue #2: It is now possible to define new handler for each language

Author: cedricrupb

code_tokenize/config.py

@@ -1,6 +1,8 @@
 
 import json
 
+from .lang.base_visitors import LeafVisitor
+
 
 class TokenizationConfig:
     """Helper object to translate arguments of tokenize to config object"""
@@ -17,7 +19,7 @@ def __init__(self, lang, **kwargs):
             "*_statement", "*_definition", "*_declaration"
         ]
 
-        self.visitors = [] # Additional visitor which should be run during analysis
+        self.visitors = [LeafVisitor] # visitor classes which should be run during analysis
 
         self.update(kwargs)
 

code_tokenize/lang/__init__.py

@@ -1,5 +1,5 @@
 
-from ..config import TokenizationConfig
+from ..config    import TokenizationConfig
 
 from .python import create_tokenization_config as pytok_config
 from .java   import create_tokenization_config as jvtok_config
@@ -8,8 +8,6 @@
 from .php    import create_tokenization_config as phptok_config
 from .ruby   import create_tokenization_config as rubytok_config
 
-from .python.indent import IndentVisitor as PythonIndentVisitor
-
 
 def load_from_lang_config(lang, **kwargs):
 
@@ -19,13 +17,7 @@ def load_from_lang_config(lang, **kwargs):
     elif lang == "javascript" : base_config = jstok_config()
     elif lang == "php"        : base_config = phptok_config()
     elif lang == "ruby"       : base_config = rubytok_config()
-    else                : base_config = TokenizationConfig(lang)
+    else                      : base_config = TokenizationConfig(lang)
 
     base_config.update(kwargs)
     return base_config
-
-
-def indent_handler_clazz(lang):
-    if lang == "python": return PythonIndentVisitor
-
-    # TODO: Ruby indent visitor

code_tokenize/lang/base_visitors.py

@@ -0,0 +1,17 @@
+from code_ast import ASTVisitor
+
+# Basic visitor -----------------------------------------------------------
+
+class LeafVisitor(ASTVisitor):
+
+    def __init__(self, node_handler):
+        self.node_handler = node_handler
+
+    def visit_string(self, node):
+        self.node_handler(node)
+        return False
+
+    def visit(self, node):
+        if node.child_count == 0:
+            self.node_handler(node)
+            return False

code_tokenize/lang/go/__init__.py

@@ -1,11 +1,30 @@
 
 from ...config import TokenizationConfig
+from ...tokens import NewlineToken
+
+from ..base_visitors import LeafVisitor
+
 
 # Tokenization config ----------------------------------------------------------------
 
 def create_tokenization_config():
     return TokenizationConfig(
         lang = 'go',
         statement_types = ["*_statement", "*_declaration"],
+        visitors = [GoLeafVisitor],
         indent_tokens   = False
-    )
+    )
+
+# Custom leaf visitor ----------------------------------------------------------------
+
+class GoLeafVisitor(LeafVisitor):
+
+    def visit_interpreted_string_literal(self, node):
+        self.node_handler(node)
+        return False
+
+    def visit(self, node):
+        if node.type == "\n":
+            self.node_handler.handle_token(NewlineToken(self.node_handler.config))
+            return False
+        return super().visit(node)

code_tokenize/lang/python/__init__.py

@@ -1,12 +1,25 @@
 
 from ...config import TokenizationConfig
 
+from ..base_visitors import LeafVisitor
+from .indent         import IndentVisitor
+
 
 # Tokenization config ----------------------------------------------------------------
 
 def create_tokenization_config():
     return TokenizationConfig(
         lang = "python",
         statement_types = ["*_statement", "*_definition"],
+        visitors = [PythonLeafVisitor, IndentVisitor],
         indent_tokens   = True
     )
+
+# Custom leaf visitor ----------------------------------------------------------------
+
+class PythonLeafVisitor(LeafVisitor):
+
+   def visit_unary_operator(self, node):
+        if node.children[-1].type == "integer":
+            self.node_handler(node)
+            return False

code_tokenize/tokenizer.py

@@ -3,8 +3,6 @@
 from code_ast.visitor import ASTVisitor, ResumingVisitorComposition
 
 from .tokens  import ASTToken, TokenSequence
-from .lang    import indent_handler_clazz
-
 
 
 # Interface ----------------------------------------------------------------
@@ -66,12 +64,14 @@ def _create_token_handler(self, code_lines):
 
     def _create_tree_visitors(self, token_handler, visitors = None):
         visitors  = visitors or []
+        visitors += self._visitor_factories
 
-        visitors += [visitor_fn(token_handler) 
-                        for visitor_fn in self._visitor_factories]
+        visitors  = [visitor_fn(token_handler) 
+                        if callable(visitor_fn) 
+                            else visitor_fn
+                        for visitor_fn in visitors]
 
         return ResumingVisitorComposition(
-            LeafVisitor(self.config, token_handler),
             ErrorVisitor(self.config),
             *visitors
         )
@@ -88,35 +88,21 @@ def __call__(self, code_tree, code_lines, visitors = None):
 
 def create_tokenizer(config):
     """Function to create tokenizer based on configuration"""
-    tokenizer = Tokenizer(config)
-
-    if config.indent_tokens:
-        indent_handler = indent_handler_clazz(config.lang)
-        assert indent_handler is not None, "Language %s does not support indentation handling" % config.lang
-        tokenizer.append_visitor(indent_handler)
-
-    return tokenizer
+    return Tokenizer(config)
 
 
 # Basic visitor -----------------------------------------------------------
 
 
 class LeafVisitor(ASTVisitor):
 
-    def __init__(self, config, node_handler):
-        self.config = config
+    def __init__(self, node_handler):
         self.node_handler = node_handler
 
     def visit_string(self, node):
         self.node_handler(node)
         return False
 
-    def visit_unary_operator(self, node):
-        # TODO Use a custom leaf visitor
-        if node.children[-1].type == "integer":
-            self.node_handler(node)
-            return False
-
     def visit(self, node):
         if node.child_count == 0:
             self.node_handler(node)

requirements.txt

@@ -1,3 +1,4 @@
 tree_sitter==0.19.0
 requests==2.25.1
-GitPython==3.1.18
+GitPython==3.1.18
+code_ast==0.1.0

tests/test_tokenization.py

@@ -42,4 +42,13 @@ def test_error_handling(self):
     def test_error_handling2(self):
         tokens = ctok.tokenize("public int myFunc(int x){\n    x = x + 1;\n    return x;\n}", lang = "java", syntax_error = "ignore")
         expected = ["public", "int", "myFunc", "", "(", "int", "x", ")", "{", "x", "=", "x", "+", "1", ";", "return", "x", ";", "}"]
+        self.assertEqual(expected, [str(t) for t in tokens])
+
+
+class GoTokenizationTest(TestCase):
+
+    def test_tokenize1(self):
+        tokens = ctok.tokenize('func main(){\n    tip1 := "test"\n}', lang = "go")
+        expected = ["func", "main", "(", ")", "{", "tip1", ":=", '"test"', "#NEWLINE#", "}"]
+
         self.assertEqual(expected, [str(t) for t in tokens])