Fixed python typecheck (#242)

src/python/engagement-creator/src/engagement-creator.py

@@ -5,7 +5,7 @@
 import time
 import traceback
 from collections import defaultdict
-from typing import Any, Dict, Iterator, Tuple, Sequence, Optional, cast
+from typing import Any, Dict, Iterator, Tuple, Sequence, Optional, cast, TypeVar
 
 import boto3
 import botocore.exceptions  # type: ignore
@@ -15,6 +15,7 @@
 from grapl_analyzerlib.prelude import BaseView
 from grapl_analyzerlib.prelude import RiskView
 from grapl_analyzerlib.viewable import Viewable
+from grapl_analyzerlib.queryable import Queryable
 from mypy_boto3_s3 import S3ServiceResource
 from mypy_boto3_sqs import SQSClient
 from typing_extensions import Type
@@ -34,6 +35,9 @@
 
 EventWithReceiptHandle = Tuple[S3Event, str]
 
+V = TypeVar("V", bound=Viewable)
+Q = TypeVar("Q", bound=Queryable)
+
 
 def parse_s3_event(s3: S3ServiceResource, event) -> bytes:
     # Retrieve body of sns message
@@ -155,10 +159,10 @@ def _upsert(client: GraphClient, node_dict: Dict[str, Any]) -> str:
 def upsert(
     client: GraphClient,
     type_name: str,
-    view_type: "Type[Viewable]",
+    view_type: Type[Viewable[V, Q]],
     node_key: str,
     node_props: Dict[str, Any],
-) -> "Viewable":
+) -> Viewable[V, Q]:
     node_props["node_key"] = node_key
     node_props["dgraph.type"] = list({type_name, "Base", "Entity"})
     uid = _upsert(client, node_props)