Repository: microsoft/pybryt

Issue:
Create a way to combine complexity classes
Add a way for users to say e.g. "linear or logarithmic complexity" in a complexity annotation. One way to do this may be to overload the `|` operator: `TimeComplexity(cplx.linear | cplx.logarithmic)`.

This would mean that the complexity classes would need to be instance of some other class instead of classes themselves. It may also be nice to define some commonly-used groups, e.g. sub-linear.

Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e3d0371..0ae0743 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,7 +3,7 @@
 All notable changes to this project will be documented in this file, and this project adheres to 
 [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## 0.4.0 - 2022-01-12
+## Unreleased
 
 * Added the `RequireImport` and `ForbidImport` annotations per [#123](https://github.com/microsoft/pybryt/issues/123)
 * Added `TimeComplexityChecker` for analyzing complexity without annotations per [#137](https://github.com/microsoft/pybryt/issues/137)
diff --git a/CITATION.cff b/CITATION.cff
index 97a782b..f2483e4 100644
--- a/CITATION.cff
+++ b/CITATION.cff
@@ -26,6 +26,6 @@ authors:
   given-names: "Nelson"
   orcid: "https://orcid.org/0000-0003-2768-5735"
 title: "Pybryt - Python library for pedagogical auto-assessment"
-version: 0.4.0
-date-released: 2022-01-12
+version: 0.3.1
+date-released: 2021-12-01
 url: "https://github.com/microsoft/pybryt"
diff --git a/README.md b/README.md
index 6b4050c..83e4034 100644
--- a/README.md
+++ b/README.md
@@ -35,7 +35,7 @@ Educators and Institutions can leverage the PyBryt Library to integrate auto ass
 
 ## Getting Started
 
-See the [Getting Started](https://microsoft.github.io/pybryt/html/getting_started.html) page on the pybryt documentation for steps to install and use pybryt for the first time. You can also check the [Microsoft Learn interactive modules](http://docs.microsoft.com/learn)on [Introductions to PyBryt](http://aka.ms/pybrytlearn) and [Advanced PyBryt](http://aka.ms/advancedpybryt) to learn more about to use the library to autoassess your learners activities.
+See the [Getting Started](https://microsoft.github.io/pybryt/html/getting_started.html) page on the pybryt documentation for steps to install and use pybryt for the first time. You can also check the [Microsoft Learn interactive module on pybryt](http://aka.ms/pybrytlearn) to learn how to use the library.
 
 ## Testing
 
diff --git a/pybryt/annotations/complexity/annotation.py b/pybryt/annotations/complexity/annotation.py
index 760e13f..e816c27 100644
--- a/pybryt/annotations/complexity/annotation.py
+++ b/pybryt/annotations/complexity/annotation.py
@@ -2,7 +2,7 @@
 
 __all__ = ["ComplexityAnnotation", "TimeComplexity"]
 
-from typing import Any, List, Tuple
+from typing import Any, List, Union
 
 from . import complexities as cplx
 
@@ -24,17 +24,30 @@ class ComplexityAnnotation(Annotation):
     from the student's memory footprint.
 
     Args:
-        complexity (:py:class:`complexity<pybryt.complexities.complexity>`): the complexity class
-            being asserted
+        complexity (:py:class:`complexity<pybryt.complexities.complexity>` or :py:class:`ComplexityUnion<pybryt.complexities.ComplexityUnion>`):
+            the complexity class being asserted
         addl_complexities(``list[complexity]``): additional custom complexity classes to consider
         **kwargs: additional keyword arguments passed to the 
             :py:class:`Annotation<pybryt.Annotation>` constructor
     """
 
-    def __init__(self, complexity: cplx.complexity, addl_complexities=[], **kwargs):
+    complexity: Union[cplx.complexity, cplx.ComplexityUnion]
+    """the complexity class of this annotation"""
+
+    addl_complexities: List[cplx.complexity]
+    """additional complexity classes to consider"""
+
+    def __init__(
+        self,
+        complexity: Union[cplx.complexity, cplx.ComplexityUnion],
+        addl_complexities: List[cplx.complexity] = [],
+        **kwargs,
+    ):
         if "name" not in kwargs:
             raise ValueError("Complexity annotations require a 'name' kwarg")
-        if complexity not in cplx.complexity_classes and not isinstance(complexity, cplx.complexity):
+        if complexity not in cplx.complexity_classes and \
+                not isinstance(complexity, cplx.complexity) and \
+                not isinstance(complexity, cplx.ComplexityUnion):
             raise ValueError(f"Invalid valid for argument 'complexity': {complexity}")
         self.complexity = complexity
         self.addl_complexities = addl_complexities
@@ -61,7 +74,7 @@ class ComplexityAnnotation(Annotation):
         Returns:
             ``bool``: whether the objects are equal
         """
-        return super().__eq__(other) and self.complexity is other.complexity
+        return super().__eq__(other) and self.complexity == other.complexity
 
 
 class TimeComplexity(ComplexityAnnotation):
@@ -95,22 +108,26 @@ class TimeComplexity(ComplexityAnnotation):
         Returns:
             :py:class:`AnnotationResult`: the results of this annotation against ``footprint``
         """
-        if self.complexity not in cplx.complexity_classes:
+        is_union = isinstance(self.complexity, cplx.ComplexityUnion)
+        if not is_union and self.complexity not in cplx.complexity_classes:
             self.addl_complexities.insert(0, self.complexity)
 
         complexity_data = {}
-        for v, ts in footprint.values:
+        for v, _ in footprint.values:
             if not isinstance(v, TimeComplexityResult) or v.name != self.name:
                 continue
 
             complexity_data[v.n] = v.stop - v.start
 
-        best_cls, best_res = None, None
-        complexities = cplx.complexity_classes + self.addl_complexities
-        for cplxy in complexities:
-            res = cplxy(complexity_data)
+        best_res = None
+        complexities: List[cplx.complexity] = cplx.complexity_classes + self.addl_complexities
+        for complexity in complexities:
+            result = complexity(complexity_data)
 
-            if best_res is None or res < best_res - EPS:
-                best_cls, best_res = cplxy, res
+            if best_res is None or result.residual < best_res.residual - EPS:
+                best_res = result
 
-        return AnnotationResult(best_cls is self.complexity, self, value=best_cls)
+        satisfied = best_res.complexity_class == self.complexity
+        if is_union:
+            satisfied = best_res.complexity_class in self.complexity.get_complexities()
+        return AnnotationResult(satisfied, self, value=best_res.complexity_class)
diff --git a/pybryt/annotations/complexity/complexities.py b/pybryt/annotations/complexity/complexities.py
index 472f92a..8c4fcc7 100644
--- a/pybryt/annotations/complexity/complexities.py
+++ b/pybryt/annotations/complexity/complexities.py
@@ -3,7 +3,21 @@
 import numpy as np
 
 from abc import ABC, abstractmethod
-from typing import Dict
+from dataclasses import dataclass
+from typing import Any, Dict, List, Union
+
+
+@dataclass
+class ComplexityClassResult:
+    """
+    A container class for the results of complexity checks.
+    """
+
+    complexity_class: "complexity"
+    """the complexity class associated with this result"""
+
+    residual: float
+    """the residual from the least-squares fit"""
 
 
 class complexity(ABC):
@@ -16,41 +30,21 @@ class complexity(ABC):
     from https://github.com/pberkes/big_O.
     """
 
-    @staticmethod
-    def __new__(cls, complexity_data: Dict[int, int]) -> float:
-        return cls.run(complexity_data)
-
-    @staticmethod
-    @abstractmethod
-    def transform_n(n: np.ndarray) -> np.ndarray:
+    def __or__(self, other: Any) -> "ComplexityUnion":
         """
-        Transforms the array of input lengths for performing least squares.
+        Create a complexity union through use of the ``|`` operator.
 
         Args:
-            n (``np.ndarray``): the array of input lengths
-        
-        Returns:
-            ``np.ndarray``: the transformed array of input lengths
-        """
-        ... # pragma: no cover
-
-    @staticmethod
-    def transform_t(t: np.ndarray) -> np.ndarray:
-        """
-        Transforms the array of timings for performing least squares.
+            other (any): the other object in the union
 
-        Args:
-            n (``np.ndarray``): the array of timings
-        
         Returns:
-            ``np.ndarray``: the transformed array of timings
+            :py:class:`ComplexityUnion`: the union
         """
-        return t
+        return ComplexityUnion.from_or(self, other)
 
-    @classmethod
-    def run(cls, complexity_data: Dict[int, int]) -> float:
+    def __call__(self, complexity_data: Dict[int, int]) -> ComplexityClassResult:
         """
-        Returns the sum of residuals by performing least squares on the input length data and timings.
+        Return the sum of residuals by performing least squares on the input length data and timings.
 
         Uses ``transform_n`` and ``transform_t`` to transform the data from ``complexity_data``, which
         is a dictionary mapping input lengths to step counts. Performs least squares on the resulting
@@ -71,18 +65,45 @@ class complexity(ABC):
         t = np.array(ts, dtype=int)
 
         try:
-            n = cls.transform_n(n)
-            t = cls.transform_t(t)
+            n = self.transform_n(n)
+            t = self.transform_t(t)
 
             _, resid, _, _ = np.linalg.lstsq(n, t, rcond=-1)
             if len(resid) == 0:
-                return np.inf
-            return resid[0]
+                return ComplexityClassResult(self, np.inf)
+            return ComplexityClassResult(self, resid[0])
         except:
-            return np.inf
+            return ComplexityClassResult(self, np.inf)
 
+    @staticmethod
+    @abstractmethod
+    def transform_n(n: np.ndarray) -> np.ndarray:
+        """
+        Transforms the array of input lengths for performing least squares.
 
-class constant(complexity):
+        Args:
+            n (``np.ndarray``): the array of input lengths
+        
+        Returns:
+            ``np.ndarray``: the transformed array of input lengths
+        """
+        ... # pragma: no cover
+
+    @staticmethod
+    def transform_t(t: np.ndarray) -> np.ndarray:
+        """
+        Transforms the array of timings for performing least squares.
+
+        Args:
+            n (``np.ndarray``): the array of timings
+        
+        Returns:
+            ``np.ndarray``: the transformed array of timings
+        """
+        return t
+
+
+class _constant(complexity):
     """
     Complexity class for constant time: :math:`\\mathcal{O}(1)`
     """
@@ -92,7 +113,11 @@ class constant(complexity):
         return np.ones((len(n), 1))
 
 
-class logarithmic(complexity):
+# complexity classes are singletons
+constant = _constant()
+
+
+class _logarithmic(complexity):
     """
     Complexity class for logarithmic time: :math:`\\mathcal{O}(\\log n)`
     """
@@ -102,7 +127,10 @@ class logarithmic(complexity):
         return np.vstack((np.ones(len(n)), np.log2(n))).T
 
 
-class linear(complexity):
+logarithmic = _logarithmic()
+
+
+class _linear(complexity):
     """
     Complexity class for linear time: :math:`\\mathcal{O}(n)`
     """
@@ -112,7 +140,10 @@ class linear(complexity):
         return np.vstack((np.ones(len(n)), n)).T
 
 
-class linearithmic(complexity):
+linear = _linear()
+
+
+class _linearithmic(complexity):
     """
     Complexity class for linearithmic time: :math:`\\mathcal{O}(n \\log n)`
     """
@@ -122,7 +153,10 @@ class linearithmic(complexity):
         return np.vstack((np.ones(len(n)), n * np.log2(n))).T
 
 
-class quadratic(complexity):
+linearithmic = _linearithmic()
+
+
+class _quadratic(complexity):
     """
     Complexity class for quadratic time: :math:`\\mathcal{O}(n^2)`
     """
@@ -132,7 +166,10 @@ class quadratic(complexity):
         return np.vstack((np.ones(len(n)), n * n)).T
 
 
-class cubic(complexity):
+quadratic = _quadratic()
+
+
+class _cubic(complexity):
     """
     Complexity class for cubic time: :math:`\\mathcal{O}(n^3)`
     """
@@ -142,7 +179,10 @@ class cubic(complexity):
         return np.vstack((np.ones(len(n)), n * n * n)).T
 
 
-class exponential(complexity):
+cubic = _cubic()
+
+
+class _exponential(complexity):
     """
     Complexity class for exponential time: :math:`\\mathcal{O}(2^n)`
     """
@@ -156,4 +196,107 @@ class exponential(complexity):
         return np.log2(t)
 
 
+exponential = _exponential()
+
+
 complexity_classes = [constant, logarithmic, linear, linearithmic, quadratic, cubic]#, exponential]
+
+
+class ComplexityUnion:
+    """
+    A complexity class that represents a union of multiple complexity classes.
+
+    Complexity unions can be used to write annotations but cannot be used as classes to determine
+    the complexity of a block of code; they are simply collections of acceptable complexity classes.
+
+    This class does not ensure the uniqueness of the complexity classes it tracks; that is, if the
+    same complexity class is added twice, it will be tracked twice.
+
+    Args:
+        *complexity_classes (:py:class:`complexity`): the complexity classes in the union
+    """
+
+    _complexity_classes = List[complexity]
+    """the complexity classes contained in the union"""
+
+    def __init__(self, *complexity_classes: complexity):
+        self._complexity_classes = list(complexity_classes)
+
+    def __or__(self, other: Any) -> "ComplexityUnion":
+        """
+        Create a complexity union through use of the ``|`` operator.
+
+        Args:
+            other (any): the other object in the union
+
+        Returns:
+            :py:class:`ComplexityUnion`: the union
+        """
+        return self.from_or(self, other)
+
+    @classmethod
+    def from_or(
+        cls,
+        left: Union[complexity, "ComplexityUnion"],
+        right: Union[complexity, "ComplexityUnion"],
+    ) -> "ComplexityUnion":
+        """
+        Create a complexity union from two operands.
+
+        Args:
+            left (:py:class:`complexity` or :py:class:`ComplexityUnion`): the left operand
+            right (:py:class:`complexity` or :py:class:`ComplexityUnion`): the right operand
+
+        Returns:
+            :py:class:`ComplexityUnion`: the union containing both operands
+
+        Raises:
+            ``TypeError``: if either operand is not a complexity class or complexity union
+        """
+        if not isinstance(left, (complexity, cls)):
+            raise TypeError(f"Attempted to combine a complexity class with an object of type {type(left)}")
+
+        if not isinstance(right, (complexity, cls)):
+            raise TypeError(f"Attempted to combine a complexity class with an object of type {type(right)}")
+
+        if isinstance(left, cls):
+            right_complexities = [right]
+            if isinstance(right, cls):
+                right_complexities = right.get_complexities()
+            return cls(*left.get_complexities(), *right_complexities)
+
+        elif isinstance(right, cls):
+            return cls(left, *right.get_complexities())
+
+        return cls(left, right)
+
+    def add_complexity(self, complexity_class: complexity) -> None:
+        """
+        Add another complexity class to this union.
+
+        Args:
+            complexity_class (:py:class:`complexity`): the complexity class to add
+        """
+        self._complexity_classes.append(complexity_class)
+
+    def get_complexities(self) -> List[complexity]:
+        """
+        Return a list of the complexity classes in this union.
+
+        Returns:
+            ``list[complexity]``: the complexity classes
+        """
+        return [*self._complexity_classes]
+
+    def __eq__(self, other: Any) -> bool:
+        """
+        Determine whether another object is equal to this complexity union.
+
+        An object is equal to a complexity union if it is also a complexity union and contains the
+        same complexity classes.
+
+        Returns:
+            ``bool``: whether the other object is equal to this one
+        """
+        return isinstance(other, type(self)) and \
+            set(self.get_complexities()) == set(other.get_complexities())
diff --git a/pybryt/version.py b/pybryt/version.py
index 6a9beea..260c070 100644
--- a/pybryt/version.py
+++ b/pybryt/version.py
@@ -1,1 +1,1 @@
-__version__ = "0.4.0"
+__version__ = "0.3.1"


Test patch:
diff --git a/tests/annotations/test_complexity.py b/tests/annotations/test_complexity.py
index 1a6d277..8b75080 100644
--- a/tests/annotations/test_complexity.py
+++ b/tests/annotations/test_complexity.py
@@ -1,6 +1,8 @@
 """Tests for complexity annotations"""
 
+from typing import Type
 import numpy as np
+from numpy.lib.arraysetops import isin
 import pytest
 
 import pybryt
@@ -84,3 +86,49 @@ def test_time_complexity():
 def test_alias():
     from pybryt.annotations.complexity import complexities as cplx2
     assert cplx.complexity_classes is cplx2.complexity_classes
+
+
+def test_complexity_union():
+    """
+    Tests for complexity unions.
+    """
+    union = cplx.logarithmic | cplx.linear
+    assert isinstance(union, cplx.ComplexityUnion)
+    assert set(union.get_complexities()) == {cplx.logarithmic, cplx.linear}
+
+    union2 = union | cplx.constant
+    assert isinstance(union2, cplx.ComplexityUnion)
+    assert set(union2.get_complexities()) == {cplx.logarithmic, cplx.linear, cplx.constant}
+
+    union3 = cplx.constant | union
+    assert isinstance(union3, cplx.ComplexityUnion)
+    assert set(union3.get_complexities()) == {cplx.logarithmic, cplx.linear, cplx.constant}
+
+    union4 = union2 | union
+    assert isinstance(union4, cplx.ComplexityUnion)
+    assert set(union4.get_complexities()) == {cplx.logarithmic, cplx.linear, cplx.constant}
+
+    union.add_complexity(cplx.linearithmic)
+    assert set(union.get_complexities()) == {cplx.logarithmic, cplx.linear, cplx.linearithmic}
+
+    assert union == cplx.logarithmic | cplx.linear | cplx.linearithmic
+    assert union != cplx.logarithmic | cplx.linear
+
+    # test whether the union works with annotations
+    a = pybryt.TimeComplexity(cplx.constant | cplx.logarithmic, name="foo")
+    footprint = generate_complexity_footprint("foo", lambda v: 1012)
+    res = a.check(footprint)
+    assert res.satisfied
+    assert res.value == cplx.constant
+
+    footprint = generate_complexity_footprint("foo", lambda v: v)
+    res = a.check(footprint)
+    assert not res.satisfied
+    assert res.value == cplx.linear
+
+    # test errors
+    with pytest.raises(TypeError):
+        cplx.constant | 1
+
+    with pytest.raises(TypeError):
+        cplx.ComplexityUnion.from_or(1, cplx.constant)