Allow disable cache

.github/actions/pytest/action.yml

@@ -17,9 +17,9 @@ runs:
       run: |
         flit install -s
         if [[ "${{ inputs.backend }}" != "" ]]; then
-          pytest --cov=datajudge --cov-report=xml --cov-append --backend=${{ inputs.backend }} ${{ inputs.args }}
+          pytest --verbose --cov=datajudge --cov-report=xml --cov-append --backend=${{ inputs.backend }} ${{ inputs.args }}
         else
-          pytest --cov=datajudge --cov-report=xml --cov-append ${{ inputs.args }}
+          pytest --verbose --cov=datajudge --cov-report=xml --cov-append ${{ inputs.args }}
         fi
     - name: Generate code coverage report
       uses: codecov/codecov-action@v3.1.3

pixi.toml

@@ -88,6 +88,10 @@ sqlalchemy = "2.*"
 pytest-cov = "*"
 pytest-xdist = "*"
 
+[feature.test.target.unix.dependencies]
+pytest-memray = "*"
+memray = "*"
+
 [feature.mypy.dependencies]
 mypy = "*"
 types-setuptools = "*"

src/datajudge/constraints/base.py

@@ -111,6 +111,10 @@ class Constraint(abc.ABC):
     In order to obtain such values, the `retrieve` method defines a mapping from DataReference,
     be it the DataReference of primary interest, `ref`, or a baseline DataReference, `ref2`, to
     value. If `ref_value` is already provided, usually no further mapping needs to be taken care of.
+
+    By default, retrieved arguments are cached indefinitely ``@lru_cache(maxsize=None)``.
+    This can be controlled by setting the `cache_size` argument to a different value.
+    ``0`` disables caching.
     """
 
     def __init__(
@@ -123,6 +127,7 @@ def __init__(
         output_processors: Optional[
             Union[OutputProcessor, List[OutputProcessor]]
         ] = output_processor_limit,
+        cache_size=None,
     ):
         self._check_if_valid_between_or_within(ref2, ref_value)
         self.ref = ref
@@ -140,6 +145,16 @@ def __init__(
             output_processors = [output_processors]
         self.output_processors = output_processors
 
+        self.cache_size = cache_size
+        self._setup_caching()
+
+    def _setup_caching(self):
+        # this has an added benefit of allowing the class to be garbage collected
+        # according to https://rednafi.com/python/lru_cache_on_methods/
+        # and https://docs.astral.sh/ruff/rules/cached-instance-method/
+        self.get_factual_value = lru_cache(self.cache_size)(self.get_factual_value)  # type: ignore[method-assign]
+        self.get_target_value = lru_cache(self.cache_size)(self.get_target_value)  # type: ignore[method-assign]
+
     def _check_if_valid_between_or_within(
         self, ref2: Optional[DataReference], ref_value: Optional[Any]
     ):
@@ -156,13 +171,13 @@ def _check_if_valid_between_or_within(
                 f"{class_name}. Use exactly either of them."
             )
 
-    @lru_cache(maxsize=None)
+    # @lru_cache(maxsize=None), see _setup_caching()
     def get_factual_value(self, engine: sa.engine.Engine) -> Any:
         factual_value, factual_selections = self.retrieve(engine, self.ref)
         self.factual_selections = factual_selections
         return factual_value
 
-    @lru_cache(maxsize=None)
+    # @lru_cache(maxsize=None), see _setup_caching()
     def get_target_value(self, engine: sa.engine.Engine) -> Any:
         if self.ref2 is None:
             return self.ref_value

src/datajudge/constraints/column.py

@@ -22,9 +22,13 @@ def retrieve(
 
 class ColumnExistence(Column):
     def __init__(
-        self, ref: DataReference, columns: List[str], name: Optional[str] = None
+        self,
+        ref: DataReference,
+        columns: List[str],
+        name: Optional[str] = None,
+        cache_size=None,
     ):
-        super().__init__(ref, ref_value=columns, name=name)
+        super().__init__(ref, ref_value=columns, name=name, cache_size=cache_size)
 
     def compare(
         self, column_names_factual: List[str], column_names_target: List[str]
@@ -86,8 +90,15 @@ def __init__(
         ref2: Optional[DataReference] = None,
         column_type: Optional[Union[str, sa.types.TypeEngine]] = None,
         name: Optional[str] = None,
+        cache_size=None,
     ):
-        super().__init__(ref, ref2=ref2, ref_value=column_type, name=name)
+        super().__init__(
+            ref,
+            ref2=ref2,
+            ref_value=column_type,
+            name=name,
+            cache_size=cache_size,
+        )
         self.column_type = column_type
 
     def retrieve(

src/datajudge/constraints/date.py

@@ -39,6 +39,7 @@ def __init__(
         use_lower_bound_reference: bool,
         column_type: str,
         name: Optional[str] = None,
+        cache_size=None,
         *,
         ref2: Optional[DataReference] = None,
         min_value: Optional[str] = None,
@@ -48,7 +49,13 @@ def __init__(
         min_date: Optional[dt.date] = None
         if min_value is not None:
             min_date = dt.datetime.strptime(min_value, INPUT_DATE_FORMAT).date()
-        super().__init__(ref, ref2=ref2, ref_value=min_date, name=name)
+        super().__init__(
+            ref,
+            ref2=ref2,
+            ref_value=min_date,
+            name=name,
+            cache_size=cache_size,
+        )
 
     def retrieve(
         self, engine: sa.engine.Engine, ref: DataReference
@@ -85,6 +92,7 @@ def __init__(
         use_upper_bound_reference: bool,
         column_type: str,
         name: Optional[str] = None,
+        cache_size=None,
         *,
         ref2: Optional[DataReference] = None,
         max_value: Optional[str] = None,
@@ -94,7 +102,13 @@ def __init__(
         max_date: Optional[dt.date] = None
         if max_value is not None:
             max_date = dt.datetime.strptime(max_value, INPUT_DATE_FORMAT).date()
-        super().__init__(ref, ref2=ref2, ref_value=max_date, name=name)
+        super().__init__(
+            ref,
+            ref2=ref2,
+            ref_value=max_date,
+            name=name,
+            cache_size=cache_size,
+        )
 
     def retrieve(
         self, engine: sa.engine.Engine, ref: DataReference
@@ -133,8 +147,9 @@ def __init__(
         lower_bound: str,
         upper_bound: str,
         name: Optional[str] = None,
+        cache_size=None,
     ):
-        super().__init__(ref, ref_value=min_fraction, name=name)
+        super().__init__(ref, ref_value=min_fraction, name=name, cache_size=cache_size)
         self.lower_bound = lower_bound
         self.upper_bound = upper_bound
 

src/datajudge/constraints/groupby.py

@@ -14,6 +14,7 @@ def __init__(
         aggregation_column: str,
         start_value: int = 0,
         name: Optional[str] = None,
+        cache_size=None,
         *,
         tolerance: float = 0,
         ref2: Optional[DataReference] = None,

src/datajudge/constraints/interval.py

@@ -19,6 +19,7 @@ def __init__(
         end_columns: List[str],
         max_relative_n_violations: float,
         name: Optional[str] = None,
+        cache_size=None,
     ):
         super().__init__(ref, ref_value=object(), name=name)
         self.key_columns = key_columns
@@ -74,6 +75,7 @@ def __init__(
         max_relative_n_violations: float,
         end_included: bool,
         name: Optional[str] = None,
+        cache_size=None,
     ):
         self.end_included = end_included
         super().__init__(
@@ -83,6 +85,7 @@ def __init__(
             end_columns,
             max_relative_n_violations,
             name=name,
+            cache_size=cache_size,
         )
 
     def select(self, engine: sa.engine.Engine, ref: DataReference):
@@ -113,6 +116,7 @@ def __init__(
         max_relative_n_violations: float,
         legitimate_gap_size: float,
         name: Optional[str] = None,
+        cache_size=None,
     ):
         self.legitimate_gap_size = legitimate_gap_size
         super().__init__(
@@ -122,6 +126,7 @@ def __init__(
             end_columns,
             max_relative_n_violations,
             name=name,
+            cache_size=cache_size,
         )
 
     @abc.abstractmethod

src/datajudge/constraints/miscs.py

@@ -9,7 +9,13 @@
 
 
 class PrimaryKeyDefinition(Constraint):
-    def __init__(self, ref, primary_keys: List[str], name: Optional[str] = None):
+    def __init__(
+        self,
+        ref,
+        primary_keys: List[str],
+        name: Optional[str] = None,
+        cache_size=None,
+    ):
         super().__init__(ref, ref_value=set(primary_keys), name=name)
 
     def retrieve(
@@ -56,6 +62,7 @@ def __init__(
         max_absolute_n_duplicates: int = 0,
         infer_pk_columns: bool = False,
         name: Optional[str] = None,
+        cache_size=None,
     ):
         if max_duplicate_fraction != 0 and max_absolute_n_duplicates != 0:
             raise ValueError(
@@ -71,7 +78,7 @@ def __init__(
             ref_value = ("relative", 0)
 
         self.infer_pk_columns = infer_pk_columns
-        super().__init__(ref, ref_value=ref_value, name=name)
+        super().__init__(ref, ref_value=ref_value, name=name, cache_size=cache_size)
 
     def test(self, engine: sa.engine.Engine) -> TestResult:
         if self.infer_pk_columns and db_access.is_bigquery(engine):
@@ -152,8 +159,15 @@ def __init__(
         max_null_fraction: Optional[float] = None,
         max_relative_deviation: float = 0,
         name: Optional[str] = None,
+        cache_size=None,
     ):
-        super().__init__(ref, ref2=ref2, ref_value=max_null_fraction, name=name)
+        super().__init__(
+            ref,
+            ref2=ref2,
+            ref_value=max_null_fraction,
+            name=name,
+            cache_size=cache_size,
+        )
         if max_null_fraction is not None and not (0 <= max_null_fraction <= 1):
             raise ValueError(
                 f"max_null_fraction was expected to lie within [0, 1] but is "

src/datajudge/constraints/nrows.py

@@ -17,8 +17,15 @@ def __init__(
         ref2: Optional[DataReference] = None,
         n_rows: Optional[int] = None,
         name: Optional[str] = None,
+        cache_size=None,
     ):
-        super().__init__(ref, ref2=ref2, ref_value=n_rows, name=name)
+        super().__init__(
+            ref,
+            ref2=ref2,
+            ref_value=n_rows,
+            name=name,
+            cache_size=cache_size,
+        )
 
     def retrieve(
         self, engine: sa.engine.Engine, ref: DataReference
@@ -85,8 +92,9 @@ def __init__(
         ref2: DataReference,
         max_relative_loss_getter: ToleranceGetter,
         name: Optional[str] = None,
+        cache_size=None,
     ):
-        super().__init__(ref, ref2=ref2, name=name)
+        super().__init__(ref, ref2=ref2, name=name, cache_size=cache_size)
         self.max_relative_loss_getter = max_relative_loss_getter
 
     def compare(self, n_rows_factual: int, n_rows_target: int) -> Tuple[bool, str]:
@@ -116,8 +124,9 @@ def __init__(
         ref2: DataReference,
         max_relative_gain_getter: ToleranceGetter,
         name: Optional[str] = None,
+        cache_size=None,
     ):
-        super().__init__(ref, ref2=ref2, name=name)
+        super().__init__(ref, ref2=ref2, name=name, cache_size=cache_size)
         self.max_relative_gain_getter = max_relative_gain_getter
 
     def compare(self, n_rows_factual: int, n_rows_target: int) -> Tuple[bool, str]:
@@ -147,8 +156,9 @@ def __init__(
         ref2: DataReference,
         min_relative_gain_getter: ToleranceGetter,
         name: Optional[str] = None,
+        cache_size=None,
     ):
-        super().__init__(ref, ref2=ref2, name=name)
+        super().__init__(ref, ref2=ref2, name=name, cache_size=cache_size)
         self.min_relative_gain_getter = min_relative_gain_getter
 
     def compare(self, n_rows_factual: int, n_rows_target: int) -> Tuple[bool, str]:

src/datajudge/constraints/numeric.py

@@ -13,11 +13,18 @@ def __init__(
         self,
         ref: DataReference,
         name: Optional[str] = None,
+        cache_size=None,
         *,
         ref2: Optional[DataReference] = None,
         min_value: Optional[float] = None,
     ):
-        super().__init__(ref, ref2=ref2, ref_value=min_value, name=name)
+        super().__init__(
+            ref,
+            ref2=ref2,
+            ref_value=min_value,
+            name=name,
+            cache_size=cache_size,
+        )
 
     def retrieve(
         self, engine: sa.engine.Engine, ref: DataReference
@@ -46,6 +53,7 @@ def __init__(
         self,
         ref: DataReference,
         name: Optional[str] = None,
+        cache_size=None,
         *,
         ref2: Optional[DataReference] = None,
         max_value: Optional[float] = None,
@@ -55,6 +63,7 @@ def __init__(
             ref2=ref2,
             ref_value=max_value,
             name=name,
+            cache_size=cache_size,
         )
 
     def retrieve(
@@ -87,8 +96,9 @@ def __init__(
         lower_bound: float,
         upper_bound: float,
         name: Optional[str] = None,
+        cache_size=None,
     ):
-        super().__init__(ref, ref_value=min_fraction, name=name)
+        super().__init__(ref, ref_value=min_fraction, name=name, cache_size=cache_size)
         self.lower_bound = lower_bound
         self.upper_bound = upper_bound
 
@@ -123,6 +133,7 @@ def __init__(
         ref: DataReference,
         max_absolute_deviation: float,
         name: Optional[str] = None,
+        cache_size=None,
         *,
         ref2: Optional[DataReference] = None,
         mean_value: Optional[float] = None,
@@ -132,6 +143,7 @@ def __init__(
             ref2=ref2,
             ref_value=mean_value,
             name=name,
+            cache_size=cache_size,
         )
         self.max_absolute_deviation = max_absolute_deviation
 
@@ -169,6 +181,7 @@ def __init__(
         max_absolute_deviation: Optional[float] = None,
         max_relative_deviation: Optional[float] = None,
         name: Optional[str] = None,
+        cache_size=None,
         *,
         ref2: Optional[DataReference] = None,
         expected_percentile: Optional[float] = None,
@@ -178,6 +191,7 @@ def __init__(
             ref2=ref2,
             ref_value=expected_percentile,
             name=name,
+            cache_size=cache_size,
         )
         if not (0 <= percentage <= 100):
             raise ValueError(