Repository: mggg/GerryChain

Issue:
Should Assignment implement collections.abc.Mapping?
Creating this in case we want to continue with the conversation started in #280 .

Abstractly, I feel like the answer is yes, but in practice this might mean adding a bunch of methods that we'll never use...

Comments:

Patch:
diff --git a/.travis.yml b/.travis.yml
deleted file mode 100644
index 12d5575..0000000
--- a/.travis.yml
+++ /dev/null
@@ -1,30 +0,0 @@
-# Config file for automatic testing at travis-ci.org
-
-language: python
-  - "2.7"
-  - "3.6"
-sudo: required
-before_install:
-  - sudo apt-get update
-install:
-  - if [[ "$TRAVIS_PYTHON_VERSION" == "2.7" ]]; then
-    wget https://repo.continuum.io/miniconda/Miniconda2-latest-Linux-x86_64.sh -O miniconda.sh;
-    else
-    wget https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh;
-    fi
-  - bash miniconda.sh -b -p $HOME/miniconda
-  - export PATH="$HOME/miniconda/bin:$PATH"
-  - hash -r
-  - conda config --set always_yes yes
-  - conda update -q conda
-  - conda install -n conda-build
-  - conda install anaconda-client
-  - conda install pytest-cov
-  - pip install coveralls
-  - conda build .
-  - conda install --use-local gerrychain
-  - conda info -a
-script:
-  - pytest -v --cov=gerrychain tests
-after_success:
-  - coveralls
diff --git a/gerrychain/partition/assignment.py b/gerrychain/partition/assignment.py
index e1c4a43..fd4f967 100644
--- a/gerrychain/partition/assignment.py
+++ b/gerrychain/partition/assignment.py
@@ -1,11 +1,13 @@
 from collections import defaultdict
+from itertools import chain, repeat
+from collections.abc import Mapping
 
 import pandas
 
 from ..updaters.flows import flows_from_changes
 
 
-class Assignment:
+class Assignment(Mapping):
     """An assignment of nodes into parts.
 
     The goal of :class:`Assignment` is to provide an interface that mirrors a
@@ -16,20 +18,21 @@ class Assignment:
     ``{part: <frozenset of nodes in part>}``.
     """
 
-    def __init__(self, parts: dict):
+    def __init__(self, parts, validate=True):
+        if validate:
+            number_of_keys = sum(len(keys) for keys in parts.values())
+            number_of_unique_keys = len(set().union(*parts.values()))
+            if number_of_keys != number_of_unique_keys:
+                raise ValueError("Keys must have unique assignments.")
+            if not all(isinstance(keys, frozenset) for keys in parts.values()):
+                raise TypeError("Level sets must be frozensets")
         self.parts = parts
 
-    @classmethod
-    def from_dict(cls, assignment):
-        """Create an Assignment from a dictionary. This is probably the method you want
-        to use to create a new assignment.
+    def __iter__(self):
+        return self.keys()
 
-        This also works for pandas Series.
-        """
-        parts = {
-            part: frozenset(nodes) for part, nodes in level_sets(assignment).items()
-        }
-        return cls(parts)
+    def __len__(self):
+        return sum(len(keys) for keys in self.parts.values())
 
     def __getitem__(self, node):
         for part, nodes in self.parts.items():
@@ -41,9 +44,9 @@ class Assignment:
         """Returns a copy of the assignment.
         Does not duplicate the frozensets of nodes, just the parts dictionary.
         """
-        return Assignment(self.parts.copy())
+        return Assignment(self.parts.copy(), validate=False)
 
-    def update(self, mapping: dict):
+    def update(self, mapping):
         """Update the assignment for some nodes using the given mapping.
         """
         flows = flows_from_changes(self, mapping)
@@ -59,6 +62,14 @@ class Assignment:
             for node in nodes:
                 yield (node, part)
 
+    def keys(self):
+        return chain(*self.parts.values())
+
+    def values(self):
+        return chain(
+            *(repeat(value, times=len(keys)) for value, keys in self.parts.items())
+        )
+
     def update_parts(self, new_parts):
         """Update some parts of the assignment. Does not check that every node is
         still assigned to a part.
@@ -69,12 +80,6 @@ class Assignment:
         for part, nodes in new_parts.items():
             self.parts[part] = frozenset(nodes)
 
-    def get(self, key, default=None):
-        try:
-            return self[key]
-        except KeyError:
-            return default
-
     def to_series(self):
         """Convert the assignment to a :class:`pandas.Series`."""
         groups = [
@@ -83,10 +88,21 @@ class Assignment:
         return pandas.concat(groups)
 
     def to_dict(self):
-        """Convert the assignment to a {node: part} dictionary.
+        """Convert the assignment to a ``{node: part}`` dictionary.
         This is expensive and should be used rarely."""
         return {node: part for part, nodes in self.parts.items() for node in nodes}
 
+    @classmethod
+    def from_dict(cls, assignment):
+        """Create an :class:`Assignment` from a dictionary. This is probably the method you want
+        to use to create a new assignment.
+
+        This also works for :class:`pandas.Series`.
+        """
+        parts = {part: frozenset(keys) for part, keys in level_sets(assignment).items()}
+
+        return cls(parts)
+
 
 def get_assignment(assignment, graph=None):
     if isinstance(assignment, str):
diff --git a/gerrychain/partition/partition.py b/gerrychain/partition/partition.py
index b9a9e5d..f263e5e 100644
--- a/gerrychain/partition/partition.py
+++ b/gerrychain/partition/partition.py
@@ -41,8 +41,8 @@ class Partition:
 
         self.assignment = get_assignment(assignment, graph)
 
-        if not self.validate_assignment():
-            raise NameError("Graph's nodes' names do not match the Assignment's geo units' names.")
+        if set(self.assignment) != set(graph):
+            raise KeyError("The graph's node labels do not match the Assignment's keys")
 
         if updaters is None:
             updaters = dict()
@@ -54,14 +54,6 @@ class Partition:
         self.flows = None
         self.edge_flows = None
 
-    def validate_assignment(self):
-        node_names = set(self.graph.nodes)
-        if len(node_names) != sum(len(dist) for dist in self.assignment.parts.values()):
-            return False
-
-        assgn_names = set(name for dist in self.assignment.parts.values() for name in dist)
-        return node_names == assgn_names
-
     def _from_parent(self, parent, flips):
         self.parent = parent
         self.flips = flips


Test patch:
diff --git a/tests/partition/test_assignment.py b/tests/partition/test_assignment.py
index 334cbe1..bdbe7e5 100644
--- a/tests/partition/test_assignment.py
+++ b/tests/partition/test_assignment.py
@@ -2,6 +2,7 @@ import pandas
 import pytest
 
 from gerrychain.partition.assignment import Assignment, get_assignment
+from collections.abc import Mapping
 
 
 @pytest.fixture
@@ -43,6 +44,50 @@ class TestAssignment:
         assignment.update_parts({2: {2}, 3: {3}})
         assert assignment.to_dict() == {1: 1, 2: 2, 3: 3}
 
+    def test_implements_Mapping_abc(self, assignment):
+        # __iter__
+        assert list(assignment) == [1, 2, 3]
+
+        # __contains__
+        for i in [1, 2, 3]:
+            assert i in assignment
+
+        # __len__
+        assert len(assignment) == 3
+
+        # __getitem__
+        assert assignment[1] == 1
+        assert assignment[3] == 2
+
+        # keys()
+        keys = list(assignment.keys())
+        assert len(keys) == 3
+        assert set(keys) == {1, 2, 3}
+
+        # values()
+        values = list(assignment.values())
+        assert len(values) == 3
+        assert set(values) == {1, 2}
+
+        # items()
+        items = list(assignment.items())
+        assert len(items) == 3
+        assert set(items) == {(1, 1), (2, 2), (3, 2)}
+
+        # __eq__
+        assert assignment == {1: 1, 2: 2, 3: 2}
+
+        assert isinstance(assignment, Mapping)
+
+    def test_assignment_raises_if_a_key_has_two_assignments(self):
+        with pytest.raises(ValueError):
+            Assignment({"one": {1, 2, 3}, "two": {1, 4, 5}})
+
+    def test_assignment_can_be_instantiated_from_series(self):
+        series = pandas.Series([1, 2, 1, 2], index=[1, 2, 3, 4])
+        assignment = Assignment.from_dict(series)
+        assert assignment == {1: 1, 2: 2, 3: 1, 4: 2}
+
 
 def test_get_assignment_accepts_assignment(assignment):
     created = assignment
@@ -53,3 +98,10 @@ def test_get_assignment_accepts_assignment(assignment):
 def test_get_assignment_raises_typeerror_for_unexpected_input():
     with pytest.raises(TypeError):
         get_assignment(None)
+
+
+def test_get_assignment_with_series():
+    series = pandas.Series([1, 2, 1, 2], index=[1, 2, 3, 4])
+    assignment = get_assignment(series)
+    assert isinstance(assignment, Assignment)
+    assert assignment == {1: 1, 2: 2, 3: 1, 4: 2}
diff --git a/tests/partition/test_partition.py b/tests/partition/test_partition.py
index ddfb2c8..4695043 100644
--- a/tests/partition/test_partition.py
+++ b/tests/partition/test_partition.py
@@ -24,22 +24,18 @@ def test_Partition_can_be_flipped(example_partition):
     assert new_partition.assignment[1] == 2
 
 
-def test_Partition_misnamed_vertices_raises_namerror():
+def test_Partition_misnamed_vertices_raises_keyerror():
     graph = networkx.complete_graph(3)
-    assignment = {'0': 1, '1': 1, '2': 2}
-    with pytest.raises(NameError):
-        partition = Partition(graph, assignment, {"cut_edges": cut_edges})
+    assignment = {"0": 1, "1": 1, "2": 2}
+    with pytest.raises(KeyError):
+        Partition(graph, assignment, {"cut_edges": cut_edges})
 
-def test_Partition_unlabelled_vertices_raises_namerror():
+
+def test_Partition_unlabelled_vertices_raises_keyerror():
     graph = networkx.complete_graph(3)
     assignment = {0: 1, 2: 2}
-    with pytest.raises(NameError):
-        partition = Partition(graph, assignment, {"cut_edges": cut_edges})
-
-
-def test_Partition_validate_vertex_in_unique_district(example_partition):
-    example_partition.assignment.parts[1] = frozenset([0,1,2])
-    assert example_partition.validate_assignment() == False
+    with pytest.raises(KeyError):
+        Partition(graph, assignment, {"cut_edges": cut_edges})
 
 
 def test_Partition_knows_cut_edges_K3(example_partition):