Addon pipeline to collect tree-sitter symbols (#1181)

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

CHANGELOG.rst

@@ -19,6 +19,10 @@ v34.5.0 (unreleased)
   SBOMs that contains properties with no values.
   https://github.com/nexB/scancode.io/issues/1185
 
+- Add a new `CollectTreeSitterSymbolsAndStrings` pipeline (addon) for collecting source
+  symbol and string using tree-sitter.
+  https://github.com/nexB/scancode.io/pull/1181
+
 v34.4.0 (2024-04-22)
 --------------------
 

docs/built-in-pipelines.rst

@@ -66,6 +66,14 @@ Collect Codebase Symbols (addon)
     :members:
     :member-order: bysource
 
+.. _collect_tree_sitter_symbols:
+
+Collect Tree-Sitter Source Symbols (addon)
+------------------------------------------
+.. autoclass:: scanpipe.pipelines.collect_tree_sitter_symbols.CollectTreeSitterSymbolsAndStrings()
+    :members:
+    :member-order: bysource
+
 .. _pipeline_find_vulnerabilities:
 
 Find Vulnerabilities (addon)

scanpipe/pipelines/collect_tree_sitter_symbols.py

@@ -28,7 +28,7 @@ class CollectTreeSitterSymbolsAndStrings(Pipeline):
     """
     Collect codebase symbols using tree-sitter and keep them in extra data field.
 
-    Also collect strings and comments.
+    Also collect strings.
     """
 
     download_inputs = False
@@ -40,7 +40,7 @@ def steps(cls):
 
     def collect_and_store_tree_sitter_symbols_and_strings(self):
         """
-        Collect symbols, strings and comments from codebase files using tree-sitter
+        Collect symbols and strings from codebase files using tree-sitter
         and store them in the extra data field.
         """
         symbols.collect_and_store_tree_sitter_symbols_and_strings(

scanpipe/pipes/symbols.py

@@ -20,8 +20,11 @@
 # ScanCode.io is a free software code scanning tool from nexB Inc. and others.
 # Visit https://github.com/nexB/scancode.io for support and download.
 
+from django.db.models import Q
+
 from source_inspector import symbols_ctags
 from source_inspector import symbols_pygments
+from source_inspector import symbols_tree_sitter
 
 from scanpipe.pipes import LoopProgress
 
@@ -103,3 +106,44 @@ def _collect_and_store_pygments_symbols_and_strings(resource):
             "source_comments": result.get("source_comments"),
         }
     )
+
+
+def collect_and_store_tree_sitter_symbols_and_strings(project, logger=None):
+    """
+    Collect symbols from codebase files using tree-sitter and store
+    them in the extra data field.
+    """
+    project_files = project.codebaseresources.files()
+
+    language_qs = Q()
+
+    for language in symbols_tree_sitter.TS_LANGUAGE_WHEELS.keys():
+        language_qs |= Q(programming_language__iexact=language)
+
+    resources = project_files.filter(
+        is_binary=False,
+        is_archive=False,
+        is_media=False,
+    ).filter(language_qs)
+
+    resources_count = resources.count()
+
+    resource_iterator = resources.iterator(chunk_size=2000)
+    progress = LoopProgress(resources_count, logger)
+
+    for resource in progress.iter(resource_iterator):
+        _collect_and_store_tree_sitter_symbols_and_strings(resource)
+
+
+def _collect_and_store_tree_sitter_symbols_and_strings(resource):
+    """
+    Collect symbols ans string from a resource using tree-sitter and store
+    them in the extra data field.
+    """
+    result = symbols_tree_sitter.get_treesitter_symbols(resource.location)
+    resource.update_extra_data(
+        {
+            "source_symbols": result.get("source_symbols"),
+            "source_strings": result.get("source_strings"),
+        }
+    )

scanpipe/tests/data/source-inspector/test3.cpp-tree-sitter-expected.json

@@ -0,0 +1,35 @@
+{
+  "source_strings": [
+    "Testing dummy functions:\n",
+    "Addition: %d + %d = %d\n",
+    "Subtraction: %d - %d = %d\n"
+  ],
+  "source_symbols": [
+    "add",
+    "a",
+    "b",
+    "a",
+    "b",
+    "subtract",
+    "a",
+    "b",
+    "a",
+    "b",
+    "main",
+    "x",
+    "y",
+    "printf",
+    "printf",
+    "x",
+    "y",
+    "add",
+    "x",
+    "y",
+    "printf",
+    "x",
+    "y",
+    "subtract",
+    "x",
+    "y"
+  ]
+}

scanpipe/tests/pipes/test_symbols.py

@@ -79,3 +79,29 @@ def test_scanpipe_pipes_collect_and_store_pygments_symbols_and_strings(self):
             expected_extra_data = json.load(f)
 
         self.assertDictEqual(expected_extra_data, result_extra_data)
+
+    def test_scanpipe_pipes_collect_and_store_tree_sitter_symbols_and_strings(self):
+        dir = self.project1.codebase_path / "codefile"
+        dir.mkdir(parents=True)
+
+        file_location = self.data_location / "source-inspector" / "test3.cpp"
+        copy_input(file_location, dir)
+
+        pipes.collect_and_create_codebase_resources(self.project1)
+
+        symbols.collect_and_store_tree_sitter_symbols_and_strings(self.project1)
+
+        main_file = self.project1.codebaseresources.files()[0]
+
+        result_extra_data = main_file.extra_data
+
+        expected_extra_data = (
+            self.data_location
+            / "source-inspector"
+            / "test3.cpp-tree-sitter-expected.json"
+        )
+
+        with open(expected_extra_data) as f:
+            expected_extra_data = json.load(f)
+
+        self.assertDictEqual(expected_extra_data, result_extra_data)

scanpipe/tests/test_pipelines.py

@@ -1297,3 +1297,35 @@ def test_scanpipe_collect_pygments_symbols_pipeline_integration(self):
             expected_extra_data = json.load(f)
 
         self.assertDictEqual(expected_extra_data, result_extra_data)
+
+    def test_scanpipe_collect_tree_sitter_symbols_pipeline_integration(self):
+        pipeline_name = "collect_tree_sitter_symbols"
+        project1 = Project.objects.create(name="Analysis")
+
+        dir = project1.codebase_path / "codefile"
+        dir.mkdir(parents=True)
+
+        file_location = self.data_location / "source-inspector" / "test3.cpp"
+        copy_input(file_location, dir)
+
+        pipes.collect_and_create_codebase_resources(project1)
+
+        run = project1.add_pipeline(pipeline_name)
+        pipeline = run.make_pipeline_instance()
+
+        exitcode, out = pipeline.execute()
+        self.assertEqual(0, exitcode, msg=out)
+
+        main_file = project1.codebaseresources.files()[0]
+        result_extra_data = main_file.extra_data
+
+        expected_extra_data = (
+            self.data_location
+            / "source-inspector"
+            / "test3.cpp-tree-sitter-expected.json"
+        )
+
+        with open(expected_extra_data) as f:
+            expected_extra_data = json.load(f)
+
+        self.assertDictEqual(expected_extra_data, result_extra_data)

setup.cfg

@@ -135,6 +135,7 @@ scancodeio_pipelines =
     collect_pygments_symbols = scanpipe.pipelines.collect_pygments_symbols:CollectPygmentsSymbolsAndStrings
     collect_source_strings = scanpipe.pipelines.collect_source_strings:CollectSourceStrings
     collect_symbols = scanpipe.pipelines.collect_symbols:CollectSymbols
+    collect_tree_sitter_symbols = scanpipe.pipelines.collect_tree_sitter_symbols:CollectTreeSitterSymbolsAndStrings
     find_vulnerabilities = scanpipe.pipelines.find_vulnerabilities:FindVulnerabilities
     inspect_elf_binaries = scanpipe.pipelines.inspect_elf_binaries:InspectELFBinaries
     inspect_packages = scanpipe.pipelines.inspect_packages:InspectPackages