Update doc strings for updaters

Author: peterrrock2

gerrychain/updaters/compactness.py

@@ -2,9 +2,20 @@
 
 from .flows import on_flow
 from .cut_edges import on_edge_flow
+from typing import Dict, Set
 
 
-def boundary_nodes(partition, alias="boundary_nodes"):
+def boundary_nodes(partition, alias: str = "boundary_nodes") -> Set:
+    """
+    :param partition: A partition of a Graph
+    :type partition: :class:`~gerrychain.partition.Partition`
+    :param alias: The name of the attribute that the boundary nodes are
+        stored under.
+    :type alias: str
+
+    :returns: The set of nodes in the partition that are on the boundary.
+    :rtype: Set
+    """
     if partition.parent:
         return partition.parent[alias]
     return {
@@ -14,20 +25,55 @@ def boundary_nodes(partition, alias="boundary_nodes"):
     }
 
 
-def initialize_exterior_boundaries_as_a_set(partition):
+def initialize_exterior_boundaries_as_a_set(partition) -> Dict[int, Set]:
+    """
+    :param partition: A partition of a Graph
+    :type partition: :class:`~gerrychain.partition.Partition`
+
+    :returns: A dictionary mapping each part of a partition to the set of nodes
+        in that part that are on the boundary.
+    :rtype: Dict[int, Set]
+    """
     part_boundaries = collections.defaultdict(set)
     for node in partition["boundary_nodes"]:
         part_boundaries[partition.assignment.mapping[node]].add(node)
     return part_boundaries
 
 
 @on_flow(initialize_exterior_boundaries_as_a_set, alias="exterior_boundaries_as_a_set")
-def exterior_boundaries_as_a_set(partition, previous, inflow, outflow):
+def exterior_boundaries_as_a_set(partition, previous: Set, inflow: Set, outflow: Set) -> Set:
+    """
+    Updater function that responds to the flow of nodes between different partitions.
+
+    :param partition: A partition of a Graph
+    :type partition: :class:`~gerrychain.partition.Partition`
+    :param previous: The previous set of exterior boundary nodes for a
+        fixed part of the given partition.
+    :type previous: Set
+    :param inflow: The set of nodes that have flowed into the given part of the
+        partition.
+    :type inflow: Set
+    :param outflow: The set of nodes that have flowed out of the given part of the
+        partition.
+    :type outflow: Set
+
+    :returns: The new set of exterior boundary nodes for the given part of the
+        partition.
+    :rtype: Set
+    """
     graph_boundary = partition["boundary_nodes"]
     return (previous | (inflow & graph_boundary)) - outflow
 
 
-def initialize_exterior_boundaries(partition):
+def initialize_exterior_boundaries(partition) -> Dict[int, float]:
+    """
+    :param partition: A partition of a Graph
+    :type partition: :class:`~gerrychain.partition.Partition`
+
+    :returns: A dictionary mapping each part of a partition to the total
+        perimeter of the boundary nodes in that part.
+    :rtype: Dict[int, float]
+    """
     graph_boundary = partition["boundary_nodes"]
     boundaries = collections.defaultdict(lambda: 0)
     for node in graph_boundary:
@@ -37,7 +83,26 @@ def initialize_exterior_boundaries(partition):
 
 
 @on_flow(initialize_exterior_boundaries, alias="exterior_boundaries")
-def exterior_boundaries(partition, previous, inflow, outflow):
+def exterior_boundaries(partition, previous: Set, inflow: Set, outflow: Set) -> Dict:
+    """
+    Updater function that responds to the flow of nodes between different partitions.
+
+    :param partition: A partition of a Graph
+    :type partition: :class:`~gerrychain.partition.Partition`
+    :param previous: The previous set of exterior boundary nodes for a
+        fixed part of the given partition.
+    :type previous: Set
+    :param inflow: The set of nodes that have flowed into the given part of the
+        partition.
+    :type inflow: Set
+    :param outflow: The set of nodes that have flowed out of the given part of the
+        partition.
+    :type outflow: Set
+
+    :returns: A dict mapping each part of the partition to the new exterior
+        boundary of that part.
+    :rtype: Dict
+    """
     graph_boundary = partition["boundary_nodes"]
     added_perimeter = sum(
         partition.graph.nodes[node]["boundary_perim"]
@@ -51,6 +116,14 @@ def exterior_boundaries(partition, previous, inflow, outflow):
 
 
 def initialize_interior_boundaries(partition):
+    """
+    :param partition: A partition of a Graph
+    :type partition: :class:`~gerrychain.partition.Partition`
+
+    :returns: A dictionary mapping each part of a partition to the total
+        perimeter the given part shares with other parts.
+    :rtype: Dict[int, float]
+    """
     return {
         part: sum(
             partition.graph.edges[edge]["shared_perim"]
@@ -61,7 +134,27 @@ def initialize_interior_boundaries(partition):
 
 
 @on_edge_flow(initialize_interior_boundaries, alias="interior_boundaries")
-def interior_boundaries(partition, previous, new_edges, old_edges):
+def interior_boundaries(partition, previous: Set, new_edges: Set, old_edges: Set) -> Dict:
+    """
+    Updater function that responds to the flow of nodes between different partitions.
+
+    :param partition: A partition of a Graph
+    :type partition: :class:`~gerrychain.partition.Partition`
+    :param previous: The previous set of exterior boundary nodes for a
+        fixed part of the given partition.
+    :type previous: Set
+    :param new_edges: The set of edges that have flowed into the given part of the
+        partition.
+    :type new_edges: Set
+    :param old_edges: The set of edges that have flowed out of the given part of the
+        partition.
+    :type old_edges: Set
+
+
+    :returns: A dict mapping each part of the partition to the new interior
+        boundary of that part.
+    :rtype: Dict
+    """
     added_perimeter = sum(
         partition.graph.edges[edge]["shared_perim"] for edge in new_edges
     )
@@ -71,21 +164,44 @@ def interior_boundaries(partition, previous, new_edges, old_edges):
     return previous + added_perimeter - removed_perimeter
 
 
-def flips(partition):
+def flips(partition) -> Dict:
+    """
+    :returns: The flips that were made to get from the parent partition to the
+        given partition.
+    :rtype: Dict
+    """
     return partition.flips
 
 
-def perimeter_of_part(partition, part):
+def perimeter_of_part(partition, part: int) -> float:
     """
     Totals up the perimeter of the part in the partition.
-    Requires that 'boundary_perim' be a node attribute, 'shared_perim' be an edge
-    attribute, 'cut_edges' be an updater, and 'exterior_boundaries' be an updater.
+
+    .. Warning::
+
+        Requires that 'boundary_perim' be a node attribute, 'shared_perim' be an edge
+        attribute, 'cut_edges' be an updater, and 'exterior_boundaries' be an updater.
+
+    :param partition: A partition of a Graph
+    :type partition: :class:`~gerrychain.partition.Partition`
+    :param part: The id of the part of the partition whose perimeter we want to compute.
+    :type part: int
+
+    :returns: The perimeter of the desired part.
+    :rtype: float
     """
     exterior_perimeter = partition["exterior_boundaries"][part]
     interior_perimeter = partition["interior_boundaries"][part]
 
     return exterior_perimeter + interior_perimeter
 
 
-def perimeter(partition):
+def perimeter(partition) -> Dict[int, float]:
+    """
+    :param partition: A partition of a Graph
+    :type partition: :class:`~gerrychain.partition.Partition`
+
+    :returns: A dictionary mapping each part of a partition to its perimeter.
+    :rtype: Dict[int, float]
+    """
     return {part: perimeter_of_part(partition, part) for part in partition.parts}

gerrychain/updaters/county_splits.py

@@ -1,25 +1,50 @@
 import collections
 from enum import Enum
+from typing import Callable, Dict
+
 
 CountyInfo = collections.namedtuple("CountyInfo", "split nodes contains")
+"""
+A named tuple to store county split information.
+
+:param split: The county split status. Makes use of
+    :class:`.CountySplit` enum to compute.
+:type split: int
+:param nodes: The nodes that are contained in the county.
+:type nodes: List
+:param contains: The assignment IDs that are contained in the county.
+:type contains: Set
+"""
 
 
 class CountySplit(Enum):
+    """
+    Enum to track county splits in a partition.
+
+    :cvar NOT_SPLIT: The county is not split.
+    :cvar NEW_SPLIT: The county is split in the current partition.
+    :cvar OLD_SPLIT: The county is split in the parent partition.
+    """
     NOT_SPLIT = 0
     NEW_SPLIT = 1
     OLD_SPLIT = 2
 
 
-def county_splits(partition_name, county_field_name):
-    """Track county splits.
+def county_splits(partition_name: str, county_field_name: str) -> Callable:
+    """
+    Update that allows for the tracking of county splits.
 
     :param partition_name: Name that the :class:`.Partition` instance will store.
+    :type partition_name: str
     :param county_field_name: Name of county ID field on the graph.
-    :return: The tracked data is a dictionary keyed on the county ID. The
+    :type county_field_name: str
+
+    :returns: The tracked data is a dictionary keyed on the county ID. The
               stored values are tuples of the form `(split, nodes, seen)`.
               `split` is a :class:`.CountySplit` enum, `nodes` is a list of
               node IDs, and `seen` is a list of assignment IDs that are
               contained in the county.
+    :rtype: Callable
     """
 
     def _get_county_splits(partition):
@@ -28,8 +53,29 @@ def _get_county_splits(partition):
     return _get_county_splits
 
 
-def compute_county_splits(partition, county_field, partition_field):
-    """Track nodes in counties and information about their splitting."""
+def compute_county_splits(partition,
+                          county_field: str,
+                          partition_field: str
+                          ) -> Dict[str, CountyInfo]:
+    """
+    Track nodes in counties and information about their splitting.
+
+    :param partition: The partition object to compute county splits for.
+    :type partition: :class:`~gerrychain.partition.Partition`
+    :param county_field: Name of county ID field on the graph.
+    :type county_field: str
+    :param partition_field: Name of the attribute in the graph
+        that stores the partition information. The county
+        split information will be computed with respect to this
+        division of the graph.
+    :type partition_field: str
+
+    :returns: A dict containing the information on how counties changed
+        between the parent and child partitions. If there is no parent
+        partition, then only the OLD_SPLIT and NOT_SPLIT values will be
+        used.
+    :rtype: Dict[str, CountyInfo]
+    """
 
     # Create the initial county data containers.
     if not partition.parent: