def _prepare_constraints(constraints):
constraints = constraints or []
rebuild_columns = set()
transform_constraints = []
reject_sampling_constraints = []
for constraint in constraints:
if isinstance(constraint, type):
constraint = constraint().to_dict()
elif isinstance(constraint, Constraint):
constraint = constraint.to_dict()
constraint = Constraint.from_dict(constraint)
if not constraint.rebuild_columns:
reject_sampling_constraints.append(constraint)
elif rebuild_columns & set(constraint.constraint_columns):
intersecting_columns = rebuild_columns & set(
constraint.constraint_columns)
raise Exception(
'Multiple constraints will modify the same column(s): '
f'"{intersecting_columns}", which may lead to the constraint '
'being unenforceable. Please use "reject_sampling" as the '
'"handling_strategy" instead.')
else:
transform_constraints.append(constraint)
rebuild_columns.update(constraint.rebuild_columns)
return reject_sampling_constraints + transform_constraints
def _fit_transform_constraints(self, data):
for idx, constraint in enumerate(self._constraints):
if isinstance(constraint, type):
constraint = constraint().to_dict()
elif isinstance(constraint, Constraint):
constraint = constraint.to_dict()
constraint = Constraint.from_dict(constraint)
self._constraints[idx] = constraint
data = constraint.fit_transform(data)
return data
def test_from_dict_subclass(self):
"""Test the ``Constraint.from_dict`` method passing a subclass name.
If the ``constraint`` string is a subclass name, take it from the
Subclasses dict.
Input:
- constraint dict with a subclass name and args
Output:
- Instance of the subclass with the right args.
"""
# Setup
constraint_dict = {
'constraint': 'UniqueCombinations',
'columns': ['a', 'b'],
}
# Run
instance = Constraint.from_dict(constraint_dict)
# Assert
assert isinstance(instance, UniqueCombinations)
assert instance._columns == ['a', 'b']
def test_from_dict_fqn(self):
"""Test the ``Constraint.from_dict`` method passing a FQN.
If the ``constraint`` string is a FQN, import the class
before creating an instance of it.
Input:
- constraint dict with a FQN and args
Output:
- Instance of the subclass with the right args.
"""
# Setup
constraint_dict = {
'constraint': 'sdv.constraints.tabular.UniqueCombinations',
'columns': ['a', 'b'],
}
# Run
instance = Constraint.from_dict(constraint_dict)
# Assert
assert isinstance(instance, UniqueCombinations)
assert instance._columns == ['a', 'b']