Merge pull request #13 from bcgsc/improvement/lint_typing_fixes_TypedDict

Improvement/lint typing fixes typed dict

Author: elewis2

pori_python/graphkb/constants.py

@@ -1,7 +1,7 @@
 import argparse
 from typing import Dict
 
-from .types import CategoryBaseTermMapping
+from pori_python.types import CategoryBaseTermMapping
 
 DEFAULT_LIMIT = 1000
 GKB_BASE_URL = "https://graphkb-api.bcgsc.ca/api"

pori_python/graphkb/genes.py

@@ -2,6 +2,8 @@
 
 from typing import Any, Dict, List, Sequence, Set, Tuple, cast
 
+from pori_python.types import Ontology, Statement, Variant
+
 from . import GraphKBConnection
 from .constants import (
     BASE_THERAPEUTIC_TERMS,
@@ -18,7 +20,6 @@
     TUMOUR_SUPPRESSIVE,
 )
 from .match import get_equivalent_features
-from .types import Ontology, Statement, Variant
 from .util import get_rid, logger
 from .vocab import get_terms_set
 

pori_python/graphkb/match.py

@@ -4,6 +4,15 @@
 
 from typing import Dict, List, Optional, Set, Union, cast
 
+from pori_python.types import (
+    BasicPosition,
+    Ontology,
+    ParsedVariant,
+    PositionalVariant,
+    Record,
+    Variant,
+)
+
 from . import GraphKBConnection
 from .constants import (
     AMBIGUOUS_AA,
@@ -15,7 +24,6 @@
     STRUCTURAL_VARIANT_TYPES,
     VARIANT_RETURN_PROPERTIES,
 )
-from .types import BasicPosition, Ontology, ParsedVariant, PositionalVariant, Record, Variant
 from .util import (
     FeatureNotFoundError,
     convert_to_rid_list,

pori_python/graphkb/statement.py

@@ -1,8 +1,9 @@
 from typing import List, cast
 
+from pori_python.types import CategoryBaseTermMapping, Statement, Variant
+
 from . import GraphKBConnection
 from .constants import FAILED_REVIEW_STATUS, RELEVANCE_BASE_TERMS, STATEMENT_RETURN_PROPERTIES
-from .types import CategoryBaseTermMapping, Statement, Variant
 from .util import convert_to_rid_list
 from .vocab import get_terms_set
 

pori_python/graphkb/types.py

@@ -10,8 +10,9 @@
 from typing import Any, Dict, Iterable, List, Optional, Union, cast
 from urllib3.util.retry import Retry
 
+from pori_python.types import OntologyTerm, ParsedVariant, PositionalVariant, Record
+
 from .constants import DEFAULT_LIMIT, DEFAULT_URL, TYPES_TO_NOTATION, AA_3to1_MAPPING
-from .types import OntologyTerm, ParsedVariant, PositionalVariant, Record
 
 QUERY_CACHE: Dict[Any, Any] = {}
 

pori_python/graphkb/util.py

@@ -1,7 +1,8 @@
 from typing import Callable, Dict, Iterable, List, Set, cast
 
+from pori_python.types import Ontology
+
 from . import GraphKBConnection
-from .types import Ontology
 from .util import convert_to_rid_list
 
 

pori_python/graphkb/vocab.py

@@ -12,12 +12,18 @@
 from pori_python.graphkb import match as gkb_match
 from pori_python.graphkb.match import INPUT_COPY_CATEGORIES
 from pori_python.graphkb.statement import get_statements_from_variants
-from pori_python.graphkb.types import Variant
 from pori_python.graphkb.util import FeatureNotFoundError
+from pori_python.types import (
+    GkbStatement,
+    IprCopyVariant,
+    IprExprVariant,
+    IprStructuralVariant,
+    KbMatch,
+    Variant,
+)
 
 from .constants import TMB_HIGH_CATEGORY
 from .ipr import convert_statements_to_alterations
-from .types import GkbStatement, IprCopyVariant, IprExprVariant, IprStructuralVariant, KbMatch
 from .util import Hashabledict, convert_to_rid_set, logger
 
 REPORTED_COPY_VARIANTS = (INPUT_COPY_CATEGORIES.AMP, INPUT_COPY_CATEGORIES.DEEP)

pori_python/ipr/annotate.py

@@ -10,14 +10,14 @@
 from typing import Callable, Dict, Iterable, List, Set, Tuple, cast
 
 from pori_python.graphkb.match import INPUT_COPY_CATEGORIES, INPUT_EXPRESSION_CATEGORIES
-
-from .types import (
+from pori_python.types import (
     IprCopyVariant,
     IprExprVariant,
     IprFusionVariant,
     IprSmallMutationVariant,
     IprVariant,
 )
+
 from .util import hash_key, logger, pandas_falsy
 
 protein_letters_3to1.setdefault("Ter", "*")

pori_python/ipr/inputs.py

@@ -8,9 +8,16 @@
 from pori_python.graphkb import GraphKBConnection
 from pori_python.graphkb import statement as gkb_statement
 from pori_python.graphkb import vocab as gkb_vocab
+from pori_python.types import (
+    GkbStatement,
+    ImageDefinition,
+    IprFusionVariant,
+    IprGene,
+    IprVariant,
+    KbMatch,
+)
 
 from .constants import GERMLINE_BASE_TERMS, VARIANT_CLASSES
-from .types import GkbStatement, ImageDefinition, IprFusionVariant, IprGene, IprVariant, KbMatch
 from .util import find_variant, logger
 
 

pori_python/ipr/ipr.py

@@ -9,6 +9,7 @@
 
 from pori_python.graphkb import GraphKBConnection
 from pori_python.graphkb.genes import get_gene_information
+from pori_python.types import IprVariant, KbMatch
 
 from .annotate import (
     annotate_copy_variants,
@@ -36,7 +37,6 @@
 )
 from .summary import summarize
 from .therapeutic_options import create_therapeutic_options
-from .types import IprVariant, KbMatch
 from .util import LOG_LEVELS, logger, trim_empty_values
 
 CACHE_GENE_MINIMUM = 5000

pori_python/ipr/main.py

@@ -6,12 +6,11 @@
 from pori_python.graphkb import GraphKBConnection
 from pori_python.graphkb.constants import RELEVANCE_BASE_TERMS
 from pori_python.graphkb.statement import categorize_relevance
-from pori_python.graphkb.types import Ontology, Record
 from pori_python.graphkb.util import convert_to_rid_list
 from pori_python.graphkb.vocab import get_term_tree
 from pori_python.ipr.inputs import create_graphkb_sv_notation
+from pori_python.types import GkbStatement, IprVariant, KbMatch, Ontology, Record
 
-from .types import GkbStatement, IprVariant, KbMatch
 from .util import (
     convert_to_rid_set,
     generate_ontology_preference_key,

pori_python/ipr/summary.py

@@ -3,11 +3,11 @@
 """
 
 import pandas
-from typing import Dict, List, Sequence
+from typing import Any, Dict, List, Sequence
 
 from pori_python.graphkb import GraphKBConnection
+from pori_python.types import IprVariant, KbMatch
 
-from .types import IprVariant, KbMatch
 from .util import (
     create_variant_name_tuple,
     find_variant,
@@ -22,7 +22,7 @@ def create_therapeutic_options(
     """
     Generate therapeutic options summary from the list of kb-matches
     """
-    options = []
+    options: List[Dict[str, Any]] = []
     resistance_markers = get_terms_set(graphkb_conn, ["no sensitivity"])
 
     for match in kb_matches:

pori_python/ipr/therapeutic_options.py

@@ -6,10 +6,8 @@
 from typing import Any, Dict, List, Sequence, Set, Tuple, cast
 
 from pori_python.graphkb import GraphKBConnection
-from pori_python.graphkb.types import Ontology, Record
 from pori_python.graphkb.vocab import get_term_tree
-
-from .types import IprVariant
+from pori_python.types import IprVariant, Ontology, Record
 
 GENE_NEIGHBORS_MAX = 3
 
@@ -100,7 +98,7 @@ def generate_ontology_preference_key(record: Ontology, sources_sort: Dict[str, i
         record.get("deprecated", False),
         record.get("alias", False),
         bool(record.get("dependency", "")),
-        sources_sort.get(record["source"], 99999),
+        sources_sort.get(str(record.get("source")), 99999),
         record["sourceId"],
         record.get("sourceIdVersion", ""),
         record["name"],

pori_python/ipr/util.py

@@ -1,13 +1,85 @@
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Tuple, TypedDict, Union
 
-try:
-    from typing import TypedDict  # type: ignore
-except ImportError:
-    from typing_extensions import TypedDict
+# TODO: Can constants in inputs.py like COPY_REQ, SMALL_MUT_REQ, just be replaced by types?
 
-from pori_python.graphkb.types import Ontology, Record
+CategoryBaseTermMapping = List[Tuple[str, List[str]]]
+Record = TypedDict("Record", {"@rid": str, "@class": str})
+EmbeddedRecord = TypedDict("EmbeddedRecord", {"@class": str})
+RecordLink = Union[str, Record]
 
-# TODO: Can constants in inputs.py like COPY_REQ, SMALL_MUT_REQ, just be replaced by types?
+
+class Ontology(Record):
+    sourceId: str
+    name: str
+    source: RecordLink
+    displayName: str
+
+
+class OntologyTerm(Record):
+    name: Optional[str]
+    sourceId: Optional[str]
+    sourceIdVersion: Optional[str]
+    displayName: Optional[str]
+
+
+OntologyLink = Union[str, Ontology]
+
+
+class BasicPosition(EmbeddedRecord):
+    pos: int
+
+
+class CytobandPosition(EmbeddedRecord):
+    arm: str
+    majorBand: str
+    minorBand: str
+
+
+Position = Union[BasicPosition, CytobandPosition]
+
+
+class Variant(Record):
+    reference1: OntologyLink
+    reference2: Optional[OntologyLink]
+    type: OntologyLink
+    zygosity: str
+    germline: bool
+    displayName: str
+
+
+class PositionalVariant(Variant):
+    break1Start: Union[Position, CytobandPosition]
+    break1End: Optional[Union[Position, CytobandPosition]]
+    break2Start: Optional[Union[Position, CytobandPosition]]
+    break2End: Optional[Union[Position, CytobandPosition]]
+    refSeq: Optional[str]
+    untemplatedSeq: Optional[str]
+    untemplatedSeqSize: Optional[int]
+
+
+class ParsedVariant(TypedDict):
+    reference1: str
+    reference2: Optional[str]
+    type: str
+    zygosity: str
+    germline: bool
+    break1Start: Union[Position, CytobandPosition]
+    break1End: Optional[Union[Position, CytobandPosition]]
+    break2Start: Optional[Union[Position, CytobandPosition]]
+    break2End: Optional[Union[Position, CytobandPosition]]
+    refSeq: Optional[str]
+    untemplatedSeq: Optional[str]
+    untemplatedSeqSize: Optional[int]
+
+
+class Statement(Record):
+    relevance: OntologyLink
+    subject: OntologyLink
+    conditions: List[OntologyLink]
+    evidence: List[OntologyLink]
+    evidenceLevel: List[OntologyLink]
+    source: RecordLink
+    sourceId: str
 
 
 class KbMatch(TypedDict):