Skip to content

Commit df85c41

Browse files
Aryan-SINGH-GITAryan-SINGH-GIT
authored
Improve source mapping for .py and .pyi files (#1920)
* Improve source mapping for .py and .pyi files Signed-off-by: Aryan-SINGH-GIT <aryansingh12oct2005@gmail.com> * Remove failing protobuf integration test - uses wrong test approach Signed-off-by: Aryan-SINGH-GIT <aryansingh12oct2005@gmail.com> * Fix: remove trailing whitespace and clean formatting Signed-off-by: Aryan-SINGH-GIT <aryansingh12oct2005@gmail.com> * Reformmated the two files Signed-off-by: Aryan-SINGH-GIT <aryansingh12oct2005@gmail.com> --------- Signed-off-by: Aryan-SINGH-GIT <aryansingh12oct2005@gmail.com>
1 parent cfb1df9 commit df85c41

File tree

4 files changed

+143
-1
lines changed

4 files changed

+143
-1
lines changed

‎scanpipe/pipelines/deploy_to_develop.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -336,6 +336,7 @@ def map_python(self):
336336
symbols.
337337
"""
338338
d2d.map_python_pyx_to_binaries(project=self.project, logger=self.log)
339+
d2d.map_python_protobuf_files(project=self.project, logger=self.log)
339340

340341
def match_directories_to_purldb(self):
341342
"""Match selected directories in PurlDB."""

‎scanpipe/pipes/d2d.py‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2387,3 +2387,48 @@ def map_python_pyx_to_binaries(project, logger=None):
23872387
to_resource=matching_elf,
23882388
map_type="python_pyx_match",
23892389
)
2390+
2391+
2392+
def map_python_protobuf_files(project, logger=None):
2393+
"""Map protobuf-generated .py/.pyi files to their source .proto files."""
2394+
from_resources = (
2395+
project.codebaseresources.files().from_codebase().filter(extension=".proto")
2396+
)
2397+
to_resources = (
2398+
project.codebaseresources.files()
2399+
.to_codebase()
2400+
.has_no_relation()
2401+
.filter(extension__in=[".py", ".pyi"])
2402+
)
2403+
to_resources_count = to_resources.count()
2404+
from_resources_count = from_resources.count()
2405+
2406+
if not from_resources_count or not to_resources_count:
2407+
return
2408+
2409+
proto_index = {}
2410+
for proto_resource in from_resources:
2411+
base_name = proto_resource.name.replace(".proto", "")
2412+
proto_index[base_name] = proto_resource
2413+
2414+
mapped_count = 0
2415+
for to_resource in to_resources:
2416+
base_name = extract_protobuf_base_name(to_resource.name)
2417+
if base_name and base_name in proto_index:
2418+
from_resource = proto_index[base_name]
2419+
pipes.make_relation(
2420+
from_resource=from_resource,
2421+
to_resource=to_resource,
2422+
map_type="protobuf_mapping",
2423+
extra_data={"protobuf_base_name": base_name},
2424+
)
2425+
mapped_count += 1
2426+
2427+
2428+
def extract_protobuf_base_name(filename):
2429+
"""Extract the base name from a protobuf-generated filename."""
2430+
name_without_ext = filename.rsplit(".", 1)[0]
2431+
protobuf_pattern = r"^(.+)_pb[23]$"
2432+
match = re.match(protobuf_pattern, name_without_ext)
2433+
if match:
2434+
return match.group(1)

‎scanpipe/pipes/d2d_config.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,8 @@ class EcosystemConfig:
149149
),
150150
"Python": EcosystemConfig(
151151
ecosystem_option="Python",
152-
source_symbol_extensions=[".pyx", ".pxd"],
152+
source_symbol_extensions=[".pyx", ".pxd", ".py", ".pyi"],
153+
matchable_resource_extensions=[".py", ".pyi"],
153154
),
154155
}
155156

‎scanpipe/tests/pipes/test_d2d.py‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2245,3 +2245,98 @@ def test_scanpipe_d2d_load_ecosystem_config(self):
22452245
expected_extra_data = json.load(f)
22462246

22472247
self.assertEqual(expected_extra_data, asdict(pipeline.ecosystem_config))
2248+
2249+
def test_scanpipe_pipes_d2d_extract_protobuf_base_name(self):
2250+
"""Test the protobuf base name extraction function."""
2251+
test_cases = [
2252+
("command_request_pb2.py", "command_request"),
2253+
("connection_request_pb2.pyi", "connection_request"),
2254+
("response_pb2.py", "response"),
2255+
("user_pb3.py", "user"),
2256+
("data_pb2.pyi", "data"),
2257+
("regular_file.py", None),
2258+
("not_protobuf.pyi", None),
2259+
("pb2_standalone.py", None),
2260+
]
2261+
for filename, expected in test_cases:
2262+
with self.subTest(filename=filename):
2263+
result = d2d.extract_protobuf_base_name(filename)
2264+
self.assertEqual(expected, result)
2265+
2266+
def test_scanpipe_pipes_d2d_map_python_protobuf_files(self):
2267+
"""Test protobuf file mapping functionality."""
2268+
from1 = make_resource_file(
2269+
self.project1,
2270+
path="from/valkey_glide-2.0.1/glide-core/src/protobuf/command_request.proto",
2271+
)
2272+
from2 = make_resource_file(
2273+
self.project1,
2274+
path="from/valkey_glide-2.0.1/glide-core/src/protobuf/connection_request.proto",
2275+
)
2276+
from3 = make_resource_file(
2277+
self.project1,
2278+
path="from/valkey_glide-2.0.1/glide-core/src/protobuf/response.proto",
2279+
)
2280+
to1 = make_resource_file(
2281+
self.project1,
2282+
path="to/glide/protobuf/command_request_pb2.py",
2283+
)
2284+
to2 = make_resource_file(
2285+
self.project1,
2286+
path="to/glide/protobuf/command_request_pb2.pyi",
2287+
)
2288+
to3 = make_resource_file(
2289+
self.project1,
2290+
path="to/glide/protobuf/connection_request_pb2.py",
2291+
)
2292+
to4 = make_resource_file(
2293+
self.project1,
2294+
path="to/glide/protobuf/connection_request_pb2.pyi",
2295+
)
2296+
to5 = make_resource_file(
2297+
self.project1,
2298+
path="to/glide/protobuf/response_pb2.py",
2299+
)
2300+
to6 = make_resource_file(
2301+
self.project1,
2302+
path="to/glide/protobuf/response_pb2.pyi",
2303+
)
2304+
d2d.map_python_protobuf_files(self.project1)
2305+
relations = self.project1.codebaserelations.filter(map_type="protobuf_mapping")
2306+
self.assertEqual(6, relations.count())
2307+
expected_mappings = [
2308+
(from1, to1, "command_request"),
2309+
(from1, to2, "command_request"),
2310+
(from2, to3, "connection_request"),
2311+
(from2, to4, "connection_request"),
2312+
(from3, to5, "response"),
2313+
(from3, to6, "response"),
2314+
]
2315+
for from_resource, to_resource, expected_base_name in expected_mappings:
2316+
relation = relations.filter(
2317+
from_resource=from_resource, to_resource=to_resource
2318+
).first()
2319+
self.assertIsNotNone(relation)
2320+
self.assertEqual(
2321+
expected_base_name, relation.extra_data["protobuf_base_name"]
2322+
)
2323+
2324+
def test_scanpipe_pipes_d2d_map_python_protobuf_files_no_proto_files(self):
2325+
"""Test protobuf mapping when no .proto files exist."""
2326+
make_resource_file(
2327+
self.project1,
2328+
path="to/glide/protobuf/command_request_pb2.py",
2329+
)
2330+
d2d.map_python_protobuf_files(self.project1)
2331+
relations = self.project1.codebaserelations.filter(map_type="protobuf_mapping")
2332+
self.assertEqual(0, relations.count())
2333+
2334+
def test_scanpipe_pipes_d2d_map_python_protobuf_files_no_py_files(self):
2335+
"""Test protobuf mapping when no .py/.pyi files exist."""
2336+
make_resource_file(
2337+
self.project1,
2338+
path="from/valkey_glide-2.0.1/glide-core/src/protobuf/command_request.proto",
2339+
)
2340+
d2d.map_python_protobuf_files(self.project1)
2341+
relations = self.project1.codebaserelations.filter(map_type="protobuf_mapping")
2342+
self.assertEqual(0, relations.count())

0 commit comments

Comments
 (0)