FEAT: implemented Session.union, .difference and .intersection (closes #22)

gdementen · gdementen · commit ddfd1d229df7 · 2021-10-11T10:22:05.000+02:00
diff --git a/doc/source/api.rst b/doc/source/api.rst
@@ -762,6 +762,16 @@ Exploring
    Session.items
    Session.summary
 
+Combining
+---------
+
+.. autosummary::
+   :toctree: _generated/
+
+   Session.union()
+   Session.intersection()
+   Session.difference()
+
 Copying
 -------
 
diff --git a/doc/source/changes/version_0_34.rst.inc b/doc/source/changes/version_0_34.rst.inc
@@ -18,6 +18,9 @@ Backward incompatible changes
 New features
 ^^^^^^^^^^^^
 
+* implemented :py:obj:`Session.union()`, :py:obj:`Session.intersection()` and :py:obj:`Session.difference()` to combine
+  several sessions into one using set operations (closes :issue:`22`).
+
 * added a feature (see the :ref:`miscellaneous section <misc>` for details). It works on :ref:`api-axis` and
   :ref:`api-group` objects.
 
diff --git a/larray/core/session.py b/larray/core/session.py
@@ -1388,6 +1388,93 @@ def apply(self, func, *args, **kwargs) -> 'Session':
         kind = kwargs.pop('kind', Array)
         return Session([(k, func(v, *args, **kwargs) if isinstance(v, kind) else v) for k, v in self.items()])
 
+    def union(self, other):
+        """Returns session with the (set) union of this session objects with other objects.
+
+        In other words, the new session will contain objects from this session then those of other whose name is not
+        present in this session. Objects relative order will be kept intact, but only unique names will be kept.
+        Objects from this session will be before objects from other.
+
+        Parameters
+        ----------
+        other : Session
+            other session
+
+        Returns
+        -------
+        Session
+
+        Examples
+        --------
+        >>> arr1, arr2, arr3 = ndtest((2, 2)), ndtest(3), ndtest((3, 2))
+        >>> s1 = Session(arr1=arr1, arr2=arr2, arr3=arr3)
+        >>> s2 = Session(arr2=arr2 + 1, arr4=arr3 + 1)
+        >>> s1.union(s2)
+        Session(arr1, arr2, arr3, arr4)
+        >>> s2.union(s1)
+        Session(arr2, arr4, arr1, arr3)
+        """
+        self_keys = set(self.keys())
+        res = self.copy()
+        res._update_from_iterable([(k, v) for k, v in other.items() if k not in self_keys])
+        return res
+
+    def intersection(self, other):
+        """Returns session with the (set) intersection of this session objects with other objects.
+
+        In other words, the new session will contain objects from this session if there is an object with the same
+        name in other. Objects relative order will be kept intact, but only unique names will be kept.
+
+        Parameters
+        ----------
+        other : Session or sequence of names
+            other names
+
+        Returns
+        -------
+        Session
+
+        Examples
+        --------
+        >>> arr1, arr2, arr3 = ndtest((2, 2)), ndtest(3), ndtest((3, 2))
+        >>> s1 = Session(arr1=arr1, arr2=arr2, arr3=arr3)
+        >>> s2 = Session(arr2=arr2 + 1, arr3=arr3 + 1)
+        >>> s1.intersection(s2)
+        Session(arr2, arr3)
+        >>> s1.intersection(['arr2', 'arr1'])
+        Session(arr1, arr2)
+        """
+        other_keys = set(other.keys() if isinstance(other, Session) else other)
+        return self[[k for k in self.keys() if k in other_keys]]
+
+    def difference(self, other):
+        """Returns session with the (set) difference of this session objects and other objects.
+
+        In other words, the new session will contain objects from this session if there is no object with the same
+        name in other. Objects relative order will be kept intact, but only unique names will be kept.
+
+        Parameters
+        ----------
+        other : Session or sequence of names
+            other names
+
+        Returns
+        -------
+        Session
+
+        Examples
+        --------
+        >>> arr1, arr2, arr3 = ndtest((2, 2)), ndtest(3), ndtest((3, 2))
+        >>> s1 = Session(arr1=arr1, arr2=arr2, arr3=arr3)
+        >>> s2 = Session(arr2=arr2 + 1, arr3=arr3 + 1)
+        >>> s1.difference(s2)
+        Session(arr1)
+        >>> s1.difference(['arr2', 'arr1'])
+        Session(arr3)
+        """
+        other_keys = set(other.keys() if isinstance(other, Session) else other)
+        return self[[k for k in self.keys() if k not in other_keys]]
+
     def summary(self, template=None) -> str:
         """
         Returns a summary of the content of the session.
