1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
|
from __future__ import annotations
import copy
import numpy as np
from xarray.backends.common import AbstractWritableDataStore
from xarray.core import indexing
from xarray.core.variable import Variable
class InMemoryDataStore(AbstractWritableDataStore):
"""
Stores dimensions, variables and attributes in ordered dictionaries, making
this store fast compared to stores which save to disk.
This store exists purely for internal testing purposes.
"""
def __init__(self, variables=None, attributes=None):
self._variables = {} if variables is None else variables
self._attributes = {} if attributes is None else attributes
def get_attrs(self):
return self._attributes
def get_variables(self):
res = {}
for k, v in self._variables.items():
v = v.copy(deep=True)
res[k] = v
v._data = indexing.LazilyIndexedArray(v._data)
return res
def get_dimensions(self):
return {d: s for v in self._variables.values() for d, s in v.dims.items()}
def prepare_variable(self, k, v, *args, **kwargs):
new_var = Variable(v.dims, np.empty_like(v), v.attrs)
self._variables[k] = new_var
return new_var, v.data
def set_attribute(self, k, v):
# copy to imitate writing to disk.
self._attributes[k] = copy.deepcopy(v)
def set_dimension(self, dim, length, unlimited_dims=None):
# in this model, dimensions are accounted for in the variables
pass
|
