Commit 41ea2601 authored by Simeon Doetsch's avatar Simeon Doetsch

Simulation: added DataObject attribute to make PlutoData customizable

parent 58180567
......@@ -20,6 +20,7 @@ class Simulation:
Simulation is subscriptable and iterable.
"""
supported_formats = ('dbl', 'flt')
DataObject = PlutoData
def __init__(self, sim_dir: str='', format: str=None, coordinates: str=None) -> None:
self.sim_dir = sim_dir
......@@ -137,7 +138,7 @@ class Simulation:
raise IndexError("Data index out of range")
return key
def __getitem__(self, key: int) -> PlutoData:
def __getitem__(self, key: int) -> self.DataObject:
"""
Access individual data frames, return them as PlutoData
If file is already loaded, object is returned, otherwise data is loaded and returned
......@@ -151,15 +152,15 @@ class Simulation:
self._data[key] = self._load_data(key)
return self._data[key]
def _load_data(self, key: int) -> PlutoData:
def _load_data(self, key: int) -> self.DataObject:
"""Load data frame"""
return PlutoData(n=self._index(key), simulation=self)
return self.DataObject(n=self._index(key), simulation=self)
def __iter__(self) -> Generator[PlutoData, None, None]:
def __iter__(self) -> Generator[self.DataObject, None, None]:
"""Iterate over all data frames"""
return self.iter()
def iter(self, start: int=0, stop: int=None, step: int=1, memory_keep: bool=False) -> Generator[PlutoData, None, None]:
def iter(self, start: int=0, stop: int=None, step: int=1, memory_keep: bool=False) -> Generator[self.DataObject, None, None]:
"""
Iterate over data frames in range, (or all).
start, stop, step [int]: works like range (start inclusive, stop exclusive)
......@@ -174,7 +175,7 @@ class Simulation:
for i in range(start, stop, step):
yield self._load_data(i)
def memory_iter(self, start=0, stop=-1, step=1) -> Generator[PlutoData, None, None]:
def memory_iter(self, start=0, stop=-1, step=1) -> Generator[self.DataObject, None, None]:
"""Deprecated, use Simulation.iter() instead"""
warnings.warn(
"plutoplot.Simulation.memory_iter() is deprecated, use plutoplot.Simulation.iter()",
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment