Repository: pangeo-forge/pangeo-forge-recipes (language: python)

Issue:
Augment Index index with start & stop 
As per discussion today, we want to do the following in `patterns.py`

```diff
@dataclass(frozen=True, order=True)
class DimIndex:
    name: str
    index: int
    sequence_len: int
    operation: CombineOp
+   start: Optional[int]
+   stop: Optional[int]
```
`start` and `stop` represent the explicit absolute array index location of the data in this chunk in the full dataset. So to insert the data, you would do (assuming 1D)
```python
target[dim_index.start:dim_index.stop] = data_from_chunk
```

We also want to write a standalone function of the form

```python
augment_index_with_start_stop(index: DimIndex, sequence_lens: List[int]) -> DimIndex:
    # take an index _without_ start / stop and add them based on the lens defined in sequence_lens
    assert len(sequence_lens) == index.sequence_len
    start = sum(sequence_lens[:index.index])
    stop = start + sequence_lens[index.index]
    # return a new dimindex because it's frozen
``` 
`sequence_lens` is the list of the lens of _all the items in the sequence_ including this one.


This can be unit tested with synthetic data.


Discussion / hints:

Fix patch:
diff --git a/pangeo_forge_recipes/patterns.py b/pangeo_forge_recipes/patterns.py
index 8364ac9d..c3c18b9b 100644
--- a/pangeo_forge_recipes/patterns.py
+++ b/pangeo_forge_recipes/patterns.py
@@ -2,7 +2,7 @@
 Filename / URL patterns.
 """
 import inspect
-from dataclasses import dataclass, field, replace
+from dataclasses import asdict, dataclass, field, replace
 from enum import Enum, auto
 from hashlib import sha256
 from itertools import product
@@ -77,12 +77,18 @@ class DimIndex:
     :param index: The position of the item within the sequence.
     :param sequence_len: The total length of the sequence.
     :param operation: What type of Combine Operation does this dimension represent.
+    :param start: The explicit absolute array index start location of the data in this chunk,
+      in respect to the full dataset.
+    :param stop: The explicit absolute array index stop location of the data in this chunk,
+      in respect to the full dataset.
     """
 
     name: str
     index: int
     sequence_len: int
     operation: CombineOp
+    start: Optional[int] = None
+    stop: Optional[int] = None
 
     def __str__(self):
         return f"{self.name}-{self.index}"
@@ -91,6 +97,9 @@ def __post_init__(self):
         assert self.sequence_len > 0
         assert self.index >= 0
         assert self.index < self.sequence_len
+        if self.start:
+            assert self.stop is not None
+            assert self.start < self.stop
 
 
 class Index(FrozenSet[DimIndex]):
@@ -101,6 +110,23 @@ class Index(FrozenSet[DimIndex]):
 FilePatternIndex = Index
 
 
+def augment_index_with_start_stop(index: DimIndex, item_lens: List[int]) -> DimIndex:
+    """Take an index _without_ start / stop and add them based on the lens defined in sequence_lens.
+
+    :param index: The ``DimIndex`` instance to augment.
+    :param item_lens: A list of integer lengths for all items in the sequence.
+    """
+
+    assert len(item_lens) == index.sequence_len
+    start = sum(item_lens[: index.index])
+    stop = start + item_lens[index.index]
+
+    kw = asdict(index)
+    kw["start"] = start
+    kw["stop"] = stop
+    return DimIndex(**kw)
+
+
 class AutoName(Enum):
     # Recommended by official Python docs for auto naming:
     # https://docs.python.org/3/library/enum.html#using-automatic-values


Test patch:
diff --git a/tests/test_patterns.py b/tests/test_patterns.py
index 9ed53ebd..62a286fe 100644
--- a/tests/test_patterns.py
+++ b/tests/test_patterns.py
@@ -5,9 +5,11 @@
 from pangeo_forge_recipes.patterns import (
     CombineOp,
     ConcatDim,
+    DimIndex,
     FilePattern,
     FileType,
     MergeDim,
+    augment_index_with_start_stop,
     pattern_from_file_sequence,
     prune_pattern,
 )
@@ -196,3 +198,27 @@ def test_setting_file_types(file_type_value):
     else:
         with pytest.raises(ValueError, match=rf"'{file_type_value}' is not a valid FileType"):
             fp = make_concat_merge_pattern(**file_type_kwargs)[0]
+
+
+@pytest.mark.parametrize(
+    "index,start,stop",
+    [(0, 0, 2), (1, 2, 4), (2, 4, 7), (3, 7, 9), (4, 9, 11)],
+)
+def test_augment_index_with_start_stop(index, start, stop):
+    kw = dict(name="time", index=index, sequence_len=5, operation=CombineOp.CONCAT)
+    di = DimIndex(**kw)
+    expected = DimIndex(start=start, stop=stop, **kw)
+    actual = augment_index_with_start_stop(di, [2, 2, 3, 2, 2])
+    assert actual == expected
+
+
+@pytest.mark.parametrize(
+    "dimindex_sequence_len, passed_sequence_lengths",
+    [(2, [1]), (2, [1, 1, 1])],
+)
+def test_augment_index_with_start_stop_raises(dimindex_sequence_len, passed_sequence_lengths):
+    kw = dict(name="time", index=0, operation=CombineOp.CONCAT)
+    di = DimIndex(sequence_len=dimindex_sequence_len, **kw)
+
+    with pytest.raises(AssertionError):
+        augment_index_with_start_stop(di, passed_sequence_lengths)