107 lines
3.3 KiB
Python
107 lines
3.3 KiB
Python
|
# mypy: ignore-errors
|
||
|
|
|
||
|
|
import functools
|
||
|
|
from typing import Optional
|
||
|
|
|
||
|
|
from .base import VariableTracker
|
||
|
|
|
||
|
|
|
||
|
|
class LazyCache:
|
||
|
|
"""Container to cache the real VariableTracker"""
|
||
|
|
|
||
|
|
def __init__(self, value, source):
|
||
|
|
assert source
|
||
|
|
self.value = value
|
||
|
|
self.source = source
|
||
|
|
self.vt: Optional[VariableTracker] = None
|
||
|
|
|
||
|
|
def realize(self, parents_tracker):
|
||
|
|
assert self.vt is None
|
||
|
|
from ..symbolic_convert import InstructionTranslator
|
||
|
|
from .builder import VariableBuilder
|
||
|
|
|
||
|
|
tx = InstructionTranslator.current_tx()
|
||
|
|
self.vt = VariableBuilder(tx, self.source)(self.value)
|
||
|
|
self.vt.parents_tracker.add(parents_tracker)
|
||
|
|
del self.value
|
||
|
|
del self.source
|
||
|
|
|
||
|
|
|
||
|
|
class LazyVariableTracker(VariableTracker):
|
||
|
|
"""
|
||
|
|
A structure that defers the creation of the actual VariableTracker
|
||
|
|
for a given underlying value until it is accessed.
|
||
|
|
|
||
|
|
The `realize` function invokes VariableBuilder to produce the real object.
|
||
|
|
Once a LazyVariableTracker has been realized, internal bookkeeping will
|
||
|
|
prevent double realization.
|
||
|
|
|
||
|
|
This object should be utilized for processing containers, or objects that
|
||
|
|
reference other objects where we may not want to take on creating all the
|
||
|
|
VariableTrackers right away.
|
||
|
|
"""
|
||
|
|
|
||
|
|
_nonvar_fields = {"_cache", *VariableTracker._nonvar_fields}
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def create(value, source, **options):
|
||
|
|
return LazyVariableTracker(LazyCache(value, source), source=source, **options)
|
||
|
|
|
||
|
|
def __init__(self, _cache, **kwargs):
|
||
|
|
assert isinstance(_cache, LazyCache)
|
||
|
|
super().__init__(**kwargs)
|
||
|
|
self._cache = _cache
|
||
|
|
|
||
|
|
def realize(self) -> VariableTracker:
|
||
|
|
"""Force construction of the real VariableTracker"""
|
||
|
|
if self._cache.vt is None:
|
||
|
|
self._cache.realize(self.parents_tracker)
|
||
|
|
return self._cache.vt
|
||
|
|
|
||
|
|
def unwrap(self):
|
||
|
|
"""Return the real VariableTracker if it already exists"""
|
||
|
|
if self.is_realized():
|
||
|
|
return self._cache.vt
|
||
|
|
return self
|
||
|
|
|
||
|
|
def is_realized(self):
|
||
|
|
return self._cache.vt is not None
|
||
|
|
|
||
|
|
def clone(self, **kwargs):
|
||
|
|
assert kwargs.get("_cache", self._cache) is self._cache
|
||
|
|
if kwargs.get("source", self.source) is not self.source:
|
||
|
|
self.realize()
|
||
|
|
return VariableTracker.clone(self.unwrap(), **kwargs)
|
||
|
|
|
||
|
|
def __str__(self):
|
||
|
|
if self.is_realized():
|
||
|
|
return self.unwrap().__str__()
|
||
|
|
return VariableTracker.__str__(self.unwrap())
|
||
|
|
|
||
|
|
def __getattr__(self, item):
|
||
|
|
return getattr(self.realize(), item)
|
||
|
|
|
||
|
|
# most methods are auto-generated below, these are the ones we want to exclude
|
||
|
|
apply = VariableTracker.apply
|
||
|
|
copy = VariableTracker.copy
|
||
|
|
__post_init__ = VariableTracker.__post_init__
|
||
|
|
__repr__ = VariableTracker.__repr__
|
||
|
|
|
||
|
|
|
||
|
|
def _create_realize_and_forward(name):
|
||
|
|
@functools.wraps(getattr(VariableTracker, name))
|
||
|
|
def realize_and_forward(self, *args, **kwargs):
|
||
|
|
return getattr(self.realize(), name)(*args, **kwargs)
|
||
|
|
|
||
|
|
return realize_and_forward
|
||
|
|
|
||
|
|
|
||
|
|
def _populate():
|
||
|
|
for name, value in VariableTracker.__dict__.items():
|
||
|
|
if name not in LazyVariableTracker.__dict__:
|
||
|
|
if callable(value):
|
||
|
|
setattr(LazyVariableTracker, name, _create_realize_and_forward(name))
|
||
|
|
|
||
|
|
|
||
|
|
_populate()
|