Improve source mapping for .py and .pyi files

Signed-off-by: Aryan-SINGH-GIT <[email protected]>

Author: Aryan-SINGH-GIT

scanpipe/pipelines/deploy_to_develop.py

@@ -295,6 +295,7 @@ def map_python(self):
         symbols.
         """
         d2d.map_python_pyx_to_binaries(project=self.project, logger=self.log)
+        d2d.map_python_protobuf_files(project=self.project, logger=self.log)
 
     def match_directories_to_purldb(self):
         """Match selected directories in PurlDB."""

scanpipe/pipes/d2d.py

@@ -2372,3 +2372,60 @@ def map_python_pyx_to_binaries(project, logger=None):
                 to_resource=matching_elf,
                 map_type="python_pyx_match",
             )
+
+
+def map_python_protobuf_files(project, logger=None):
+    """Map protobuf-generated .py/.pyi files to their source .proto files."""
+    from_resources = (
+        project.codebaseresources.files()
+        .from_codebase()
+        .filter(extension=".proto")
+    )
+    to_resources = (
+        project.codebaseresources.files()
+        .to_codebase()
+        .has_no_relation()
+        .filter(extension__in=[".py", ".pyi"])
+    )
+
+    to_resources_count = to_resources.count()
+    from_resources_count = from_resources.count()
+
+    if not from_resources_count:
+        return
+
+    if not to_resources_count:
+        return
+
+    proto_index = {}
+    for proto_resource in from_resources:
+        base_name = proto_resource.name.replace(".proto", "")
+        proto_index[base_name] = proto_resource
+
+    mapped_count = 0
+    for to_resource in to_resources:
+        base_name = _extract_protobuf_base_name(to_resource.name)
+        
+        if base_name and base_name in proto_index:
+            from_resource = proto_index[base_name]
+            pipes.make_relation(
+                from_resource=from_resource,
+                to_resource=to_resource,
+                map_type="protobuf_mapping",
+                extra_data={"protobuf_base_name": base_name},
+            )
+            mapped_count += 1
+
+
+def _extract_protobuf_base_name(filename):
+    """Extract the base name from a protobuf-generated filename."""
+    import re
+    
+    name_without_ext = filename.rsplit(".", 1)[0]
+    protobuf_pattern = r"^(.+)_pb[23]$"
+    match = re.match(protobuf_pattern, name_without_ext)
+    
+    if match:
+        return match.group(1)
+    
+    return None

scanpipe/pipes/d2d_config.py

@@ -148,7 +148,8 @@ class EcosystemConfig:
     ),
     "Python": EcosystemConfig(
         ecosystem_option="Python",
-        source_symbol_extensions=[".pyx", ".pxd"],
+        source_symbol_extensions=[".pyx", ".pxd", ".py", ".pyi"],
+        matchable_resource_extensions=[".py", ".pyi"],
     ),
 }
 

scanpipe/tests/pipes/test_d2d.py

@@ -2155,3 +2155,107 @@ def test_scanpipe_d2d_load_ecosystem_config(self):
             expected_extra_data = json.load(f)
 
         self.assertEqual(expected_extra_data, asdict(pipeline.ecosystem_config))
+
+    def test_scanpipe_pipes_d2d_extract_protobuf_base_name(self):
+        """Test the protobuf base name extraction function."""
+        test_cases = [
+            ("command_request_pb2.py", "command_request"),
+            ("connection_request_pb2.pyi", "connection_request"),
+            ("response_pb2.py", "response"),
+            ("user_pb3.py", "user"),
+            ("data_pb2.pyi", "data"),
+            ("regular_file.py", None),
+            ("not_protobuf.pyi", None),
+            ("pb2_standalone.py", None),
+        ]
+        
+        for filename, expected in test_cases:
+            with self.subTest(filename=filename):
+                result = d2d._extract_protobuf_base_name(filename)
+                self.assertEqual(expected, result)
+
+    def test_scanpipe_pipes_d2d_map_python_protobuf_files(self):
+        """Test protobuf file mapping functionality."""
+        from1 = make_resource_file(
+            self.project1,
+            path="from/valkey_glide-2.0.1/glide-core/src/protobuf/command_request.proto",
+        )
+        from2 = make_resource_file(
+            self.project1,
+            path="from/valkey_glide-2.0.1/glide-core/src/protobuf/connection_request.proto",
+        )
+        from3 = make_resource_file(
+            self.project1,
+            path="from/valkey_glide-2.0.1/glide-core/src/protobuf/response.proto",
+        )
+        
+        to1 = make_resource_file(
+            self.project1,
+            path="to/glide/protobuf/command_request_pb2.py",
+        )
+        to2 = make_resource_file(
+            self.project1,
+            path="to/glide/protobuf/command_request_pb2.pyi",
+        )
+        to3 = make_resource_file(
+            self.project1,
+            path="to/glide/protobuf/connection_request_pb2.py",
+        )
+        to4 = make_resource_file(
+            self.project1,
+            path="to/glide/protobuf/connection_request_pb2.pyi",
+        )
+        to5 = make_resource_file(
+            self.project1,
+            path="to/glide/protobuf/response_pb2.py",
+        )
+        to6 = make_resource_file(
+            self.project1,
+            path="to/glide/protobuf/response_pb2.pyi",
+        )
+        
+        d2d.map_python_protobuf_files(self.project1)
+        
+        relations = self.project1.codebaserelations.filter(map_type="protobuf_mapping")
+        self.assertEqual(6, relations.count())
+        
+        expected_mappings = [
+            (from1, to1, "command_request"),
+            (from1, to2, "command_request"),
+            (from2, to3, "connection_request"),
+            (from2, to4, "connection_request"),
+            (from3, to5, "response"),
+            (from3, to6, "response"),
+        ]
+        
+        for from_resource, to_resource, expected_base_name in expected_mappings:
+            relation = relations.filter(
+                from_resource=from_resource,
+                to_resource=to_resource
+            ).first()
+            self.assertIsNotNone(relation)
+            self.assertEqual(expected_base_name, relation.extra_data["protobuf_base_name"])
+
+    def test_scanpipe_pipes_d2d_map_python_protobuf_files_no_proto_files(self):
+        """Test protobuf mapping when no .proto files exist."""
+        make_resource_file(
+            self.project1,
+            path="to/glide/protobuf/command_request_pb2.py",
+        )
+        
+        d2d.map_python_protobuf_files(self.project1)
+        
+        relations = self.project1.codebaserelations.filter(map_type="protobuf_mapping")
+        self.assertEqual(0, relations.count())
+
+    def test_scanpipe_pipes_d2d_map_python_protobuf_files_no_py_files(self):
+        """Test protobuf mapping when no .py/.pyi files exist."""
+        make_resource_file(
+            self.project1,
+            path="from/valkey_glide-2.0.1/glide-core/src/protobuf/command_request.proto",
+        )
+        
+        d2d.map_python_protobuf_files(self.project1)
+        
+        relations = self.project1.codebaserelations.filter(map_type="protobuf_mapping")
+        self.assertEqual(0, relations.count())

scanpipe/tests/test_pipelines.py

@@ -1701,6 +1701,73 @@ def test_scanpipe_deploy_to_develop_pipeline_integration_elfs(self):
         expected_file = self.data / "d2d-elfs" / "brotli-elf-d2d.json"
         self.assertPipelineResultEqual(expected_file, result_file)
 
+    def test_scanpipe_deploy_to_develop_pipeline_integration_protobuf(self):
+        """Test the map_deploy_to_develop pipeline with protobuf files."""
+        pipeline_name = "map_deploy_to_develop"
+        project1 = make_project(name="ProtobufTest")
+        selected_groups = ["Python"]
+
+
+        from_dir = project1.codebase_path / "from"
+        from_dir.mkdir(parents=True)
+        
+
+        proto_dir = from_dir / "valkey_glide-2.0.1" / "glide-core" / "src" / "protobuf"
+        proto_dir.mkdir(parents=True)
+        
+        proto_files = [
+            "command_request.proto",
+            "connection_request.proto", 
+            "response.proto"
+        ]
+        
+        for proto_file in proto_files:
+            (proto_dir / proto_file).write_text(f'syntax = "proto3";\npackage glide;\nmessage {proto_file.replace(".proto", "").title()} {{}}')
+        
+
+        to_dir = project1.codebase_path / "to"
+        to_dir.mkdir(parents=True)
+        
+        py_dir = to_dir / "glide" / "protobuf"
+        py_dir.mkdir(parents=True)
+        
+        py_files = [
+            "command_request_pb2.py",
+            "command_request_pb2.pyi",
+            "connection_request_pb2.py",
+            "connection_request_pb2.pyi",
+            "response_pb2.py",
+            "response_pb2.pyi"
+        ]
+        
+        for py_file in py_files:
+            (py_dir / py_file).write_text(f'# Generated by protoc\nclass {py_file.replace("_pb2.py", "").replace("_pb2.pyi", "")}: pass')
+        
+
+        run = project1.add_pipeline(
+            pipeline_name=pipeline_name, selected_groups=selected_groups
+        )
+        pipeline = run.make_pipeline_instance()
+
+        exitcode, out = pipeline.execute()
+        self.assertEqual(0, exitcode, msg=out)
+
+    
+        protobuf_relations = project1.codebaserelations.filter(map_type="protobuf_mapping")
+        self.assertGreater(protobuf_relations.count(), 0)
+        
+   
+        command_request_relations = protobuf_relations.filter(
+            to_resource__path__contains="command_request_pb2"
+        )
+        self.assertGreater(command_request_relations.count(), 0)
+        
+       
+        for relation in protobuf_relations:
+            self.assertIn(".proto", relation.from_resource.path)
+            self.assertIn("_pb2", relation.to_resource.path)
+            self.assertEqual("protobuf_mapping", relation.map_type)
+
     def test_scanpipe_deploy_to_develop_pipeline_extract_input_files_errors(self):
         project1 = make_project()
         run = project1.add_pipeline("map_deploy_to_develop")