diff --git a/dfhack/python/pydfhack/construction.py b/dfhack/python/pydfhack/construction.py index d0876c5d9..c19be2fc1 100644 --- a/dfhack/python/pydfhack/construction.py +++ b/dfhack/python/pydfhack/construction.py @@ -3,23 +3,16 @@ Python class for DF_Hack::Construction """ from ._pydfhack import _ConstructionManager -class Construction(_ConstructionManager): +from .mixins import NeedsStart +from .decorators import suspend + +class Construction(NeedsStart, _ConstructionManager): api = None - started = False + cls = _ConstructionManager def __init__(self, api, *args, **kwds): - _ConstructionManager.__init__(self, args, kwds) + self.cls.__init__(self, args, kwds) self.api = api - def prepare(self): - """ - Enforce Suspend/Start - """ - if self.api.prepare(): - if not self.started: - self.started = self.Start() - return self.started - else: - return False - - - + @suspend + def Read(self, *args, **kw): + return self.cls.Read(self, *args, **kw) diff --git a/dfhack/python/pydfhack/vegetation.py b/dfhack/python/pydfhack/vegetation.py index b80291d66..ce1262a13 100644 --- a/dfhack/python/pydfhack/vegetation.py +++ b/dfhack/python/pydfhack/vegetation.py @@ -3,23 +3,16 @@ Python class for DF_Hack::Vegetation """ from ._pydfhack import _VegetationManager -class Vegetation(_VegetationManager): +from .mixins import NeedsStart +from .decorators import suspend + +class Vegetation(NeedsStart, _VegetationManager): api = None - started = False + cls = _VegetationManager def __init__(self, api, *args, **kwds): - _VegetationManager.__init__(self, args, kwds) + self.cls.__init__(self, args, kwds) self.api = api - - def prepare(self): - """ - Enforce Suspend/Start - """ - if self.api.prepare(): - if not self.started: - self.started = self.Start() - return self.started - else: - return False - - + @suspend + def Read(self, *args, **kw): + return self.cls.Read(self, *args, **kw)