Type annotations

src/glypy/structure/fragment.py

@@ -558,7 +558,7 @@ def y_fragments_from_links(links_to_break: List['Link'], **kwargs):
 
     for ids, subtree in unique_subtrees:
         subtree = subtree.reroot(index_method=None)
-        include_nodes = {n.id for n in subtree}
+        include_nodes = ids
 
         link_ids = [link.id for link in links_to_break
                     if link.parent.id in include_nodes or

src/glypy/structure/glycan.py

@@ -513,7 +513,7 @@ def depth_first_traversal(self, from_node=None, apply_fn=identity, visited=None)
     traversal_methods['depth_first_traversal'] = "depth_first_traversal"
 
     def breadth_first_traversal(self, from_node=None, apply_fn=identity, visited=None):
-        '''
+        """
         Make a breadth-first traversal of the glycan graph. Children are explored in descending bond-order.
 
         When selecting an iteration strategy, this strategy is specified as "bfs".
@@ -536,7 +536,7 @@ def breadth_first_traversal(self, from_node=None, apply_fn=identity, visited=Non
         See also
         --------
         Glycan.depth_first_traversal
-        '''
+        """
         node_queue = deque([self.root if from_node is None else from_node])
         visited = set() if visited is None else visited
         while len(node_queue) > 0:
@@ -549,8 +549,7 @@ def breadth_first_traversal(self, from_node=None, apply_fn=identity, visited=Non
                 res = apply_fn(node)
                 if res is not None:
                     yield res
-            # node_queue.extend(terminal for link in node.links.values()
-            #                   for terminal in link if terminal.id not in visited)
+
             for link in node.links.values():
                 terminal = link.parent
                 if terminal.id not in visited:
@@ -616,9 +615,11 @@ def indexed_traversal(self, from_node=None, apply_fn=identity, visited=None):
                 i += 1
 
     traversal_methods['index'] = "indexed_traversal"
+    traversal_methods['indexed_traversal'] = "indexed_traversal"
 
     def _get_traversal_method(self, method):
-        """An internal helper method used to resolve traversal
+        """
+        An internal helper method used to resolve traversal
         methods by name or alias.
 
         Parameters

src/glypy/structure/link.py

@@ -1,15 +1,15 @@
 import itertools
 from uuid import uuid4
-try:
-    from collections.abc import Iterable
-except ImportError:
-    from collections import Iterable
+
+from typing import Iterable, Optional, List, TYPE_CHECKING
 
 from glypy.composition import Composition
 from glypy.utils import uid, basestring, make_struct
 from .base import SaccharideBase, SubstituentBase
 from .constants import UnknownPosition, LinkageType
 
+if TYPE_CHECKING:
+    from .base import MoleculeBase
 
 default_parent_loss = Composition({"O": 1, "H": 1})
 default_child_loss = Composition(H=1)
@@ -41,13 +41,36 @@ class Link(object):
     '''
 
     __slots__ = (
-        "parent", "child", "parent_position", "child_position",
-        "parent_loss", "child_loss", "id",
-        "label", "_attached",
+        "parent",
+        "child",
+        "parent_position",
+        "child_position",
+        "parent_loss",
+        "child_loss",
+        "id",
+        "label",
+        "_attached",
         "parent_linkage_type",
         "child_linkage_type"
     )
 
+    parent: Optional['MoleculeBase']
+    child: Optional['MoleculeBase']
+
+    parent_position: int
+    child_position: int
+
+    parent_loss: Composition
+    child_loss: Composition
+
+    id: Optional[int]
+    label: Optional[str]
+
+    parent_linkage_type: LinkageType
+    child_linkage_type: LinkageType
+
+    _attached: Optional[bool]
+
     def __init__(self, parent, child, parent_position=UnknownPosition, child_position=UnknownPosition,
                  parent_loss=None, child_loss=None, id=None, attach=True, parent_linkage_type=None,
                  child_linkage_type=None):
@@ -524,6 +547,11 @@ class LinkMaskContext(object):
     '''
     A context manager for masking and unmasking |Link| objects on a residue
     '''
+
+    residue: "MoleculeBase"
+    links: List[Link]
+    attach: bool
+
     def __init__(self, residue, attach=False):
         self.residue = residue
         self.links = [link for link in residue.links.values()]
@@ -575,6 +603,12 @@ class AmbiguousLink(Link):
         "child_choices", "child_position_choices"
     )
 
+    parent_choices: List["MoleculeBase"]
+    child_choices: List["MoleculeBase"]
+
+    parent_position_choices: List[int]
+    child_position_choices: List[int]
+
     def __init__(self, parent, child, parent_position=(UnknownPosition,), child_position=(UnknownPosition,),
                  parent_loss=None, child_loss=None, id=None, attach=True, parent_linkage_type=None,
                  child_linkage_type=None):

src/glypy/structure/monosaccharide.py

@@ -1,4 +1,5 @@
 import logging
+from typing import Optional, Tuple
 try:
     from itertools import chain, izip_longest
 except ImportError:
@@ -336,6 +337,21 @@ class Monosaccharide(SaccharideBase):
         "_checked_for_reduction"
     )
 
+    id: int
+    _anomer: Anomer
+    _configuration: Tuple[Configuration]
+    _stem: Tuple[Stem]
+    _superclass: SuperClass
+    ring_start: Optional[int]
+    rind_end: Optional[int]
+    links: OrderedMultiMap[int, Link]
+    substituent_links: OrderedMultiMap[int, Link]
+    modifications: OrderedMultiMap[int, Modification]
+    composition: Composition
+    _reducing_end: Optional['ReducedEnd']
+    _degree: int
+    _checked_for_reduction: Optional[bool]
+
     def __init__(self, anomer=None, configuration=None, stem=None,
                  superclass=None, ring_start=UnknownPosition, ring_end=UnknownPosition,
                  modifications=None, links=None, substituent_links=None,

src/glypy/utils/base.py

@@ -13,6 +13,16 @@
 except ImportError:  # pragma: no cover
     from xml.etree import ElementTree as ET
 
+from typing import (List, Type, Optional, Protocol, TYPE_CHECKING,
+                    Iterable, TypeVar, DefaultDict, Hashable,
+                    Callable, Generic, Dict, Mapping)
+
+if TYPE_CHECKING:
+    from glypy.structure.base import MoleculeBase
+    from glypy.structure import Glycan
+
+T = TypeVar("T")
+K = TypeVar("K", bound=Hashable)
 
 i128 = int
 basestring = (bytes, str)
@@ -52,7 +62,7 @@ def opener(obj, mode='r'):
         raise IOError("Can't find a way to open {}".format(obj))
 
 
-def invert_dict(d):
+def invert_dict(d: Mapping[K, T]) -> Dict[T, K]:
     return {v: k for k, v in d.items()}
 
 
@@ -81,19 +91,19 @@ def count_up():
     return count_up
 
 
-def identity(x):   # pragma: no cover
+def identity(x: T) -> T:   # pragma: no cover
     return x
 
 
 def nullop(*args, **kwargs):   # pragma: no cover
     pass
 
 
-def chrinc(a='a', i=1):
+def chrinc(a: str='a', i: int=1) -> str:
     return chr(ord(a) + i)
 
 
-def make_struct(name, fields, debug=False, docstring=None):
+def make_struct(name: str, fields: List[str], debug: bool=False, docstring: str=Optional[None]) -> Type:
     '''
     A convenience function for defining plain-old-data (POD) objects that are optimized
     for named accessor lookup, unlike `namedtuple`. If the named container does not
@@ -179,20 +189,20 @@ def __dict__(self):
     return result
 
 
-class ClassPropertyDescriptor(object):
+class ClassPropertyDescriptor(Generic[T]):
     '''
     Standard Class Property Descriptor Implementation
     '''
     def __init__(self, fget, fset=None):
         self.fget = fget
         self.fset = fset
 
-    def __get__(self, obj, klass=None):
+    def __get__(self, obj, klass=None) -> T:
         if klass is None:  # pragma: no cover
             klass = type(obj)
         return self.fget.__get__(obj, klass)()
 
-    def __set__(self, obj, value):
+    def __set__(self, obj, value: T):
         if not self.fset:  # pragma: no cover
             raise AttributeError("can't set attribute")
         type_ = type(obj)
@@ -223,7 +233,13 @@ class RootProtocolNotSupportedError(TypeError):
     pass
 
 
-def root(structure):
+
+class Rootable(Protocol):
+    def __root__(self) -> "MoleculeBase":
+        ...
+
+
+def root(structure: Rootable):
     """A generic method for obtaining the root of a structure representing
     or containing a glycan graph with a single distinct root.
 
@@ -259,7 +275,12 @@ class TreeProtocolNotSupportedError(TypeError):
     pass
 
 
-def tree(structure):
+class Treeable(Protocol):
+    def __tree__(self) -> "Glycan":
+        ...
+
+
+def tree(structure: Treeable):
     """A generic method for obtaining the :class:`Glycan` of a structure representing
     or containing a glycan graph.
 
@@ -288,18 +309,18 @@ def tree(structure):
     return tree
 
 
-def groupby(ungrouped_list, key_fn=identity):
+def groupby(ungrouped_list: Iterable[T], key_fn: Callable[[T], K]=identity) -> DefaultDict[K, List[T]]:
     groups = defaultdict(list)
     for item in ungrouped_list:
         key_value = key_fn(item)
         groups[key_value].append(item)
     return groups
 
 
-def where(iterable, fn):
+def where(iterable: Iterable[T], fn: Callable[[T], bool]) -> List[int]:
     return [i for i, k in enumerate(iterable) if fn(k)]
 
 
-def uid(n=128):
+def uid(n=128) -> int:
     int_ = random.getrandbits(n)
     return int_