Fix type hints.

src/ostorlab/agent/mixins/agent_report_vulnerability_mixin.py

@@ -13,6 +13,19 @@
 from ostorlab.assets import asset as os_asset
 
 
+ArgProtoDicType = Dict[str, Union[str, bytes]]
+InputProtoDictType = Dict[str, Union[str, bytes]]
+FrameProtoDictType = Dict[
+    str, Union[str, int, List[ArgProtoDicType], List[InputProtoDictType]]
+]
+StackTraceProtoDictType = Dict[
+    str, Union[str, float, List[FrameProtoDictType], List[InputProtoDictType]]
+]
+VulnerabilityLocationMetadataProtoDictType = Dict[
+    str, Union[str, StackTraceProtoDictType]
+]
+
+
 class MetadataType(enum.Enum):
     """Vulnerability location metadata type."""
 
@@ -35,8 +48,8 @@ class Arg:
     value: Optional[bytes] = None
     type: Optional[str] = None
 
-    def to_proto_dict(self) -> Dict[str, Union[Any, List[Dict[str, str]]]]:
-        proto_dict_value = {}
+    def to_proto_dict(self) -> ArgProtoDicType:
+        proto_dict_value: ArgProtoDicType = {}
         if self.name is not None:
             proto_dict_value["name"] = self.name
         if self.value is not None:
@@ -59,8 +72,8 @@ class Frame:
     address: Optional[int] = None
     args: List[Arg] = dataclasses.field(default_factory=lambda: [])
 
-    def to_proto_dict(self) -> Dict[str, Union[Any, List[Dict[str, str]]]]:
-        proto_dict_value = {"function_name": self.function_name}
+    def to_proto_dict(self) -> FrameProtoDictType:
+        proto_dict_value: FrameProtoDictType = {"function_name": self.function_name}
         if self.class_name is not None:
             proto_dict_value["class_name"] = self.class_name
         if self.package_name is not None:
@@ -84,8 +97,8 @@ class Input:
     value: Optional[bytes] = None
     type: Optional[str] = None
 
-    def to_proto_dict(self) -> Dict[str, Union[Any, List[Dict[str, str]]]]:
-        proto_dict_value = {}
+    def to_proto_dict(self) -> InputProtoDictType:
+        proto_dict_value: InputProtoDictType = {}
         if self.name is not None:
             proto_dict_value["name"] = self.name
         if self.value is not None:
@@ -104,8 +117,10 @@ class StackTrace:
     time: Optional[float] = None
     inputs: List[Input] = dataclasses.field(default_factory=lambda: [])
 
-    def to_proto_dict(self) -> Dict[str, Union[Any, List[Dict[str, str]]]]:
-        proto_dict_value = {"frames": [frame.to_proto_dict() for frame in self.frames]}
+    def to_proto_dict(self) -> StackTraceProtoDictType:
+        proto_dict_value: StackTraceProtoDictType = {
+            "frames": [frame.to_proto_dict() for frame in self.frames]
+        }
         if self.thread_id is not None:
             proto_dict_value["thread_id"] = self.thread_id
         if self.time is not None:
@@ -124,10 +139,12 @@ class VulnerabilityLocationMetadata:
     and the value of 'where' the vulnerability was detected in the asset"""
 
     metadata_type: MetadataType
-    value: str | StackTrace
+    value: Union[str, StackTrace]
 
-    def to_proto_dict(self) -> Dict[str, Union[Any, List[Dict[str, str]]]]:
-        proto_dict_value = {"type": self.metadata_type.name}
+    def to_proto_dict(self) -> VulnerabilityLocationMetadataProtoDictType:
+        proto_dict_value: VulnerabilityLocationMetadataProtoDictType = {
+            "type": self.metadata_type.name
+        }
         if isinstance(self.value, str):
             proto_dict_value["raw_value"] = self.value
         if isinstance(self.value, StackTrace):
@@ -142,11 +159,13 @@ class VulnerabilityLocation:
     metadata: List[VulnerabilityLocationMetadata]
     asset: Union[os_asset.Asset, None] = None
 
-    def to_dict(self) -> Dict[str, Union[Any, List[Dict[str, str]]]]:
+    def to_dict(
+        self,
+    ) -> Dict[str, Union[Any, List[VulnerabilityLocationMetadataProtoDictType]]]:
         """Convert data class to a dict matching what is expected from protobuf."""
-        location: Dict[str, Union[Dict[str, str], List[Dict[str, str]]]] = {
-            "metadata": [meta.to_proto_dict() for meta in self.metadata]
-        }
+        location: Dict[
+            str, Union[Any, List[VulnerabilityLocationMetadataProtoDictType]]
+        ] = {"metadata": [meta.to_proto_dict() for meta in self.metadata]}
         if self.asset is not None:
             location[self.asset.proto_field] = self.asset.__dict__
         return location