Fix MyFiniteSetDisjointUnion and pass more tests

TestFinitePosetRepresentation
TestFinitePosetConstructionSum

3 files changed, 45 insertions, 15 deletions

diff --git a/src/act4e_solutions/posets_representation.py b/src/act4e_solutions/posets_representation.py
index 99fafdb..61a598c 100644
--- a/src/act4e_solutions/posets_representation.py
+++ b/src/act4e_solutions/posets_representation.py
@@ -7,6 +7,7 @@ import numpy as np
 
 from .sets_representation import SolFiniteSetRepresentation
 from .posets_product import MyFinitePosetProduct
+from .posets_sum import MyFinitePosetDisjointUnion
 
 import act4e_interfaces as I
 
@@ -119,14 +120,12 @@ class MyFinitePoset(I.FinitePoset[X]):
 class SolFinitePosetRepresentation(I.FinitePosetRepresentation):
     def load(self, h: I.IOHelper, s: I.FinitePoset_desc) -> I.FinitePoset[Any]:
         if "poset_product" in s.keys():
-            posets = []
-            for p in s["poset_product"]:
-                posets.append(self.load(h, p))
-
+            posets = [self.load(h, p) for p in s["poset_product"]]
             return MyFinitePosetProduct(posets)
 
         if "poset_sum" in s.keys():
-            raise NotImplementedError
+            posets = [self.load(h, p) for p in s["poset_sum"]]
+            return MyFinitePosetDisjointUnion(posets)
 
         if not all(k in s.keys() for k in ["carrier", "hasse"]):
             raise I.InvalidFormat()
diff --git a/src/act4e_solutions/posets_sum.py b/src/act4e_solutions/posets_sum.py
index a797a14..c26a6e5 100644
--- a/src/act4e_solutions/posets_sum.py
+++ b/src/act4e_solutions/posets_sum.py
@@ -1,11 +1,39 @@
-from typing import Any, overload, Sequence, TypeVar
+from typing import Any, overload, Sequence, TypeVar, Tuple
+
+from .sets_sum import MyFiniteSetDisjointUnion
 
 import act4e_interfaces as I
 
-X = TypeVar("X")
 
+C = TypeVar("C")
+E = Tuple[int, C]
+
+class MyFinitePosetDisjointUnion(I.FinitePosetDisjointUnion[C, E]):
+    _posets: Sequence[I.FinitePoset[C]]
+    _carrier: I.FiniteSetDisjointUnion[C, E]
+    def __init__(self, posets: Sequence[I.FinitePoset[C]]):
+        self._posets = posets
+        self._carrier = MyFiniteSetDisjointUnion(p.carrier() for p in posets)
+
+    def components(self) -> Sequence[I.FinitePoset[C]]:
+        return self._posets
+
+    def carrier(self) -> I.FiniteSetDisjointUnion[C, E]:
+        return self._carrier
+
+    def holds(self, a: E, b: E):
+        ia, va = self._carrier.unpack(a)
+        ib, vb = self._carrier.unpack(b)
+
+        if ia != ib:
+            return False
+
+        return self._posets[ia].holds(va, vb)
+
+
+X = TypeVar("X")
 
 class SolFinitePosetConstructionSum(I.FinitePosetConstructionSum):
     def disjoint_union(self, ps: Sequence[I.FinitePoset[X]]) -> I.FinitePosetDisjointUnion[X, Any]:
-        raise NotImplementedError() # implement here
+        return MyFinitePosetDisjointUnion(ps)
 
diff --git a/src/act4e_solutions/sets_sum.py b/src/act4e_solutions/sets_sum.py
index 9a3f364..f36cec1 100644
--- a/src/act4e_solutions/sets_sum.py
+++ b/src/act4e_solutions/sets_sum.py
@@ -1,4 +1,4 @@
-from typing import Any, Sequence, TypeVar, List, Tuple
+from typing import Any, Sequence, TypeVar, List, Tuple, cast
 
 import act4e_interfaces as I
 
@@ -18,16 +18,19 @@ class MyFiniteSetDisjointUnion(I.FiniteSetDisjointUnion[C, E]):
     def size(self):
         return sum(c.size() for c in self._components)
 
-    def save(self, h: I.IOHelper, x: E) -> List[E]:
-        return cast(I.ConcreteRepr, list(x))
+    def save(self, h: I.IOHelper, x: E) -> I.ConcreteRepr:
+        return cast(I.ConcreteRepr, list(self.unpack(x)))
 
     def load(self, h: I.IOHelper, o: I.ConcreteRepr):
         return cast(E, o)
 
     def elements(self) -> Sequence[E]:
+        es = []
         for i, c in enumerate(self._components):
             for e in c.elements():
-                yield [i, c]
+                es.append((i, e))
+
+        return es
 
     def components(self) -> Sequence[I.FiniteSet[C]]:
         return self._components
@@ -36,10 +39,10 @@ class MyFiniteSetDisjointUnion(I.FiniteSetDisjointUnion[C, E]):
         if i < 0 or i > len(self._components):
             raise I.InvalidValue
 
-        # if not self._components[i].contains(e):
-        #     raise I.InvalidValue
+        if not self._components[i].contains(e):
+            raise I.InvalidValue
 
-        return [i, e]
+        return (i, e)
 
     def unpack(self, e: E) -> Tuple[int, C]:
         return tuple(e)