Use a NamedTuple instead of a TypedDict for TraceFragment

Summary: Let's use a `NamedTuple` instead of a `TypedDict` for `TraceFragment`. NamedTuple has a more compact representation (it is a `tuple` under the hood) and gives stronger type guarantees (cannot add arbitrary key-value pairs). Reviewed By: tianhan0 Differential Revision: D58524515 fbshipit-source-id: 2a533174255e98e0020e64f9639012f212af6349
facebook · Jun 13, 2024 · 81c777f · 81c777f
1 parent 8c0e21d
commit 81c777f
Showing 1 changed file with 48 additions and 59 deletions.
diff --git a/sapp/pipeline/pysa_taint_parser.py b/sapp/pipeline/pysa_taint_parser.py
@@ -79,7 +79,8 @@ def discard_port(self) -> LeafWithDistance:
         return LeafWithDistance(name=self.name, kind=self.kind, distance=self.distance)
 
 
-class TraceFragment(TypedDict):
+# Represents a trace frame (source or sink) in a format similar to SAPP.
+class TraceFragment(NamedTuple):
     callee: str
     port: str
     location: ParsePosition
@@ -201,20 +202,16 @@ def _parse_model_sources(
                 yield ParseConditionTuple(
                     type=ParseType.POSTCONDITION,
                     caller=callable,
-                    callee=fragment["callee"],
-                    callee_location=SourceLocation.from_typed_dict(
-                        fragment["location"]
-                    ),
-                    filename=fragment["location"]["filename"],
-                    titos=list(map(SourceLocation.from_typed_dict, fragment["titos"])),
-                    leaves=[(leaf.kind, leaf.distance) for leaf in fragment["leaves"]],
+                    callee=fragment.callee,
+                    callee_location=SourceLocation.from_typed_dict(fragment.location),
+                    filename=fragment.location["filename"],
+                    titos=list(map(SourceLocation.from_typed_dict, fragment.titos)),
+                    leaves=[(leaf.kind, leaf.distance) for leaf in fragment.leaves],
                     caller_port=port,
-                    callee_port=fragment["port"],
-                    type_interval=fragment["type_interval"],
-                    features=flatten_features_to_parse_trace_feature(
-                        fragment["features"]
-                    ),
-                    annotations=fragment["trace_annotations"],
+                    callee_port=fragment.port,
+                    type_interval=fragment.type_interval,
+                    features=flatten_features_to_parse_trace_feature(fragment.features),
+                    annotations=fragment.trace_annotations,
                 )
 
     def _parse_model_sinks(
@@ -226,20 +223,16 @@ def _parse_model_sinks(
                 yield ParseConditionTuple(
                     type=ParseType.PRECONDITION,
                     caller=callable,
-                    callee=fragment["callee"],
-                    callee_location=SourceLocation.from_typed_dict(
-                        fragment["location"]
-                    ),
-                    filename=fragment["location"]["filename"],
-                    titos=list(map(SourceLocation.from_typed_dict, fragment["titos"])),
-                    leaves=[(leaf.kind, leaf.distance) for leaf in fragment["leaves"]],
+                    callee=fragment.callee,
+                    callee_location=SourceLocation.from_typed_dict(fragment.location),
+                    filename=fragment.location["filename"],
+                    titos=list(map(SourceLocation.from_typed_dict, fragment.titos)),
+                    leaves=[(leaf.kind, leaf.distance) for leaf in fragment.leaves],
                     caller_port=port,
-                    callee_port=fragment["port"],
-                    type_interval=fragment["type_interval"],
-                    features=flatten_features_to_parse_trace_feature(
-                        fragment["features"]
-                    ),
-                    annotations=fragment["trace_annotations"],
+                    callee_port=fragment.port,
+                    type_interval=fragment.type_interval,
+                    features=flatten_features_to_parse_trace_feature(fragment.features),
+                    annotations=fragment.trace_annotations,
                 )
 
     @log_trace_keyerror_in_generator
@@ -326,21 +319,19 @@ def _parse_issue_trace_fragments(
 
         for trace in traces:
             for fragment in self._parse_trace_fragment(leaf_port, trace):
-                leaves = fragment["leaves"]
+                leaves = fragment.leaves
                 fragments.append(
                     ParseIssueConditionTuple(
-                        callee=fragment["callee"],
-                        port=fragment["port"],
-                        location=SourceLocation.from_typed_dict(fragment["location"]),
+                        callee=fragment.callee,
+                        port=fragment.port,
+                        location=SourceLocation.from_typed_dict(fragment.location),
                         leaves=[(leaf.kind, leaf.distance) for leaf in leaves],
-                        titos=list(
-                            map(SourceLocation.from_typed_dict, fragment["titos"])
-                        ),
+                        titos=list(map(SourceLocation.from_typed_dict, fragment.titos)),
                         features=flatten_features_to_parse_trace_feature(
-                            fragment["features"]
+                            fragment.features
                         ),
-                        type_interval=fragment["type_interval"],
-                        annotations=fragment["trace_annotations"],
+                        type_interval=fragment.type_interval,
+                        annotations=fragment.trace_annotations,
                     )
                 )
                 leaf_distances.update(
@@ -390,17 +381,16 @@ def _parse_trace_fragment(
                 (callee_name, callee_port),
                 leaves,
             ) in leaf_name_and_port_to_leaves.items():
-                fragment: TraceFragment = {
-                    "callee": callee_name,
-                    "port": callee_port,
-                    "location": location,
-                    "leaves": leaves,
-                    "titos": tito_positions,
-                    "features": local_features,
-                    "type_interval": type_interval,
-                    "trace_annotations": trace_annotations,
-                }
-                yield fragment
+                yield TraceFragment(
+                    callee=callee_name,
+                    port=callee_port,
+                    location=location,
+                    leaves=leaves,
+                    titos=tito_positions,
+                    features=local_features,
+                    type_interval=type_interval,
+                    trace_annotations=trace_annotations,
+                )
         elif "call" in trace:
             call = trace["call"]
             port = call["port"]
@@ -411,17 +401,16 @@ def _parse_trace_fragment(
             ]
 
             for resolved in resolves_to:
-                fragment: TraceFragment = {
-                    "callee": resolved,
-                    "port": port,
-                    "location": location,
-                    "leaves": leaves,
-                    "titos": tito_positions,
-                    "features": local_features,
-                    "type_interval": type_interval,
-                    "trace_annotations": trace_annotations,
-                }
-                yield fragment
+                yield TraceFragment(
+                    callee=resolved,
+                    port=port,
+                    location=location,
+                    leaves=leaves,
+                    titos=tito_positions,
+                    features=local_features,
+                    type_interval=type_interval,
+                    trace_annotations=trace_annotations,
+                )
         elif "declaration" in trace:
             pass  # User-declared fragment.
         else: