| name | replace-decorator |
| description | Reference for using, reviewing, or modifying unitorch's @replace decorator in src/unitorch/utils/decorators.py. Use when overriding upstream classes, adding replacement classes under modules/replace, reasoning about process-global monkey patches, or debugging import-time replacement and subclass __bases__ rewriting. |
@replace Decorator Reference
@replace is defined in src/unitorch/utils/decorators.py.
Signature
from unitorch.utils.decorators import replace
@replace(TargetClass)
class ReplacementClass(TargetClass):
...
Mechanism
- Record
TargetClass -> ReplacementClass in the module-level
OPTIMIZED_CLASSES dict.
- Set
ReplacementClass.__replaced_class__ = TargetClass.
- Walk
sys.modules and:
- Replace every module-level name that equals
TargetClass with
ReplacementClass.
- Rewrite
__bases__ of any class that inherits from TargetClass.
When To Use
Use @replace when you need to override upstream library behavior such as
HuggingFace diffusers or datasets without forking the library or changing
call sites. The replacement class typically:
- Inherits from the target to reuse its logic.
- Overrides specific methods to fix bugs, skip validation, or add features.
Conventions In This Codebase
| Location | Pattern |
|---|
src/unitorch/modules/replace/diffusers_v2.py | Override diffusers pipeline __call__ / check_inputs. |
src/unitorch/modules/replace/datasets_v2.py | Override HuggingFace datasets iterables for fast skip support. |
Replacement classes are named <Original>V2 by convention and are decorated
immediately after their definition.
Important Constraints
- The
@replace call must happen at module import time. Do not put it inside a
function or conditional block.
- The replacement runs once when the module is first imported. Re-importing has
no additional effect, although a warning is logged if the same target is
replaced twice.
- The replacement is process-global. It affects every consumer of the patched
module in the same Python process.
- Always inherit from the target class so that
isinstance checks and
super() calls continue to work correctly.
