Add type annotations

abnumber/alignment.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import List, Union
 
 from abnumber.common import is_similar_residue, is_integer
 from abnumber.position import Position
@@ -28,7 +28,7 @@ class Alignment:
     ...
 
     """
-    def __init__(self, positions, residues, scheme, chain_type):
+    def __init__(self, positions: List, residues: List, scheme, chain_type):
         assert isinstance(positions, list), 'Expected list of positions and residues. ' \
                                             'Use chain.align(other) to create an alignment.'
         assert len(positions) == len(residues)
@@ -58,7 +58,7 @@ def __getitem__(self, item):
         raw_pos = self.positions.index(pos)
         return self.residues[raw_pos]
 
-    def slice(self, start: Union[str, int, 'Position'] = None, stop: Union[str, int, 'Position'] = None,
+    def slice(self, start: Union[str, int, 'Position', None] = None, stop: Union[str, int, 'Position', None] = None,
               stop_inclusive: bool = True, allow_raw: bool = False):
         """Create a slice of this alignment
 

abnumber/chain.py

@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from typing import Union, List, Generator, Tuple
+from typing import Union, List, Generator, Tuple, Optional, Literal
 from Bio import SeqIO
 from Bio.SeqRecord import SeqRecord
 import pandas as pd
@@ -14,6 +14,11 @@
 from abnumber.position import Position
 
 
+ValidSchemes = Literal['imgt', 'chothia', 'kabat', 'aho']
+ValidCDRDefs = Literal['imgt', 'chothia', 'kabat', 'north']
+ValidSpecies = Literal['human', 'mouse', 'rat', 'rabbit', 'rhesus', 'pig', 'alpaca']
+
+
 class Chain:
     """
     Antibody chain aligned to a chosen antibody numbering scheme
@@ -63,7 +68,7 @@ class Chain:
     :param germline: (Internal use only) Germline as identified by ANARCI
     """
 
-    def __init__(self, sequence, scheme, cdr_definition=None, name=None, assign_germline=False, allowed_species=None, **kwargs):
+    def __init__(self, sequence: str, scheme: Optional[ValidSchemes], cdr_definition: Optional[ValidCDRDefs] = None, name: Optional[str] = None, assign_germline=False, allowed_species: Optional[ValidSpecies] = None, **kwargs):
         aa_dict = kwargs.pop('aa_dict', None)
         chain_type = kwargs.pop('chain_type', None)
         tail = kwargs.pop('tail', None)

abnumber/position.py

@@ -1,5 +1,5 @@
 import copy
-from typing import List, Union
+from typing import List
 
 from abnumber.common import _validate_chain_type, SCHEME_POSITION_TO_REGION, SCHEME_VERNIER, POS_REGEX
 
@@ -147,7 +147,7 @@ def is_light_chain(self):
         return self.chain_type in 'KL'
 
 
-def sort_positions(positions: List[str], chain_type: str, scheme: str) -> List:
+def sort_positions(positions: List[str], chain_type: str, scheme: str) -> List[str]:
     """Sort position strings to correct order based on given scheme"""
     has_prefix = [p.startswith('H') or p.startswith('L') for p in positions]
     assert all(has_prefix) or not any(has_prefix), 'Inconsistent position prefix'