add sequence_len to DimIndex
MODIFY
pangeo_forge_recipes/patterns.py
pangeo_forge_recipes/patterns.py
@@ -14,7 +14,7 @@ class CombineOp(Enum):
SUBSET = 3
-@dataclass
+@dataclass(frozen=True)
class ConcatDim:
"""Represents a concatenation operation across a dimension of a FilePattern.
@@ -35,7 +35,7 @@ class ConcatDim:
operation: ClassVar[CombineOp] = CombineOp.CONCAT
-@dataclass
+@dataclass(frozen=True)
class MergeDim:
"""Represents a merge operation across a dimension of a FilePattern.
@@ -56,8 +56,14 @@ class MergeDim:
class DimIndex:
name: str
index: int
+ sequence_len: int
operation: CombineOp
+ def __post_init__(self):
+ assert self.sequence_len > 0
+ assert self.index >= 0
+ assert self.index < self.sequence_len
+
FilePatternIndex = Tuple[DimIndex, ...]
CombineDim = Union[MergeDim, ConcatDim]
@@ -136,7 +142,8 @@ class FilePattern:
"""Iterate over all keys in the pattern. """
for val in product(*[range(n) for n in self.shape]):
yield tuple(
- [DimIndex(op.name, v, op.operation) for op, v in zip(self.combine_dims, val)]
+ DimIndex(op.name, v, len(op.keys), op.operation)
+ for op, v in zip(self.combine_dims, val)
)
def items(self):

MODIFY
tests/test_patterns.py
tests/test_patterns.py
@@ -69,8 +69,10 @@ def test_file_pattern_concat_merge(pickle):
for k in key:
if k.name == "time":
assert k.operation == CombineOp.CONCAT
+ assert k.sequence_len == 3
if k.name == "variable":
assert k.operation == CombineOp.MERGE
+ assert k.sequence_len == 2
assert fp[key] == fname

add sequence_len to DimIndex