diff --git a/pyrit/common/apply_defaults.py b/pyrit/common/apply_defaults.py index 7fdb7b8b22..c27b294444 100644 --- a/pyrit/common/apply_defaults.py +++ b/pyrit/common/apply_defaults.py @@ -106,6 +106,16 @@ def set_default_value( parameter_name=parameter_name, include_subclasses=include_subclasses, ) + # A re-registration under the opposite flag replaces the previous one + # entirely; keeping both would leave the older value reachable (the + # lookup checks the include_subclasses=True scope first) and subclasses + # stuck inheriting it. + opposite_scope = DefaultValueScope( + class_type=class_type, + parameter_name=parameter_name, + include_subclasses=not include_subclasses, + ) + self._default_values.pop(opposite_scope, None) self._default_values[scope] = value logger.debug(f"Set default value for {class_type.__name__}.{parameter_name} = {value}") @@ -125,14 +135,18 @@ def get_default_value( Returns: Tuple of (found, value) where found indicates if a default was found. """ - # First, try exact match - scope = DefaultValueScope( - class_type=class_type, - parameter_name=parameter_name, - include_subclasses=True, - ) - if scope in self._default_values: - return True, self._default_values[scope] + # First, try exact match for both registration flags. A default + # registered with include_subclasses=False must still apply to the + # registered class itself - the flag only controls whether subclasses + # inherit the default. + for include_subclasses in (True, False): + scope = DefaultValueScope( + class_type=class_type, + parameter_name=parameter_name, + include_subclasses=include_subclasses, + ) + if scope in self._default_values: + return True, self._default_values[scope] # Then, check parent classes if include_subclasses is True for existing_scope, value in self._default_values.items(): diff --git a/tests/unit/common/test_apply_defaults.py b/tests/unit/common/test_apply_defaults.py index 77f472ee07..198a472cc4 100644 --- a/tests/unit/common/test_apply_defaults.py +++ b/tests/unit/common/test_apply_defaults.py @@ -98,6 +98,26 @@ def test_global_default_values_no_subclass_when_disabled(): assert found is False +def test_global_default_values_no_subclass_still_applies_to_base(): + registry = GlobalDefaultValues() + registry.set_default_value(class_type=_Base, parameter_name="name", value="no-inherit", include_subclasses=False) + found, val = registry.get_default_value(class_type=_Base, parameter_name="name") + assert found is True + assert val == "no-inherit" + child = _Child() + assert child.name is None + + +def test_global_default_values_mixed_flags_same_param(): + registry = GlobalDefaultValues() + registry.set_default_value(class_type=_Base, parameter_name="name", value="base-only", include_subclasses=False) + registry.set_default_value(class_type=_Child, parameter_name="name", value="child-default") + found, val = registry.get_default_value(class_type=_Base, parameter_name="name") + assert found is True and val == "base-only" + found, val = registry.get_default_value(class_type=_Child, parameter_name="name") + assert found is True and val == "child-default" + + def test_global_default_values_reset(): registry = GlobalDefaultValues() registry.set_default_value(class_type=_Base, parameter_name="name", value="x") diff --git a/tests/unit/common/test_pyrit_default_value.py b/tests/unit/common/test_pyrit_default_value.py index a6366ca615..40e2203b32 100644 --- a/tests/unit/common/test_pyrit_default_value.py +++ b/tests/unit/common/test_pyrit_default_value.py @@ -214,6 +214,47 @@ def __init__( assert child_obj.param2 == 50 assert child_obj.param3 == 3.14 + def test_reregistering_true_then_false_replaces_previous(self) -> None: + """True-then-False re-registration must replace the old scope, not shadow it.""" + + class TargetClass: + @apply_defaults + def __init__(self, *, param1: str | None = None) -> None: + self.param1 = param1 + + class ChildClass(TargetClass): + @apply_defaults + def __init__(self, *, param1: str | None = None) -> None: + super().__init__(param1=param1) + + set_default_value(class_type=TargetClass, parameter_name="param1", value="broad", include_subclasses=True) + set_default_value(class_type=TargetClass, parameter_name="param1", value="exact", include_subclasses=False) + + # The newer registration wins for the registered class itself... + assert TargetClass().param1 == "exact" + # ...and subclasses no longer inherit the removed True-scope default. + assert ChildClass().param1 is None + + def test_reregistering_false_then_true_replaces_previous(self) -> None: + """False-then-True re-registration restores subclass inheritance.""" + + class TargetClass: + @apply_defaults + def __init__(self, *, param1: str | None = None) -> None: + self.param1 = param1 + + class ChildClass(TargetClass): + @apply_defaults + def __init__(self, *, param1: str | None = None) -> None: + super().__init__(param1=param1) + + set_default_value(class_type=TargetClass, parameter_name="param1", value="exact", include_subclasses=False) + set_default_value(class_type=TargetClass, parameter_name="param1", value="broad", include_subclasses=True) + + assert TargetClass().param1 == "broad" + # Subclass inheritance comes back with the restored True scope. + assert ChildClass().param1 == "broad" + def test_parent_not_affected_by_child_defaults(self) -> None: """Test that setting defaults on child class doesn't affect parent instances."""