Check various keys for netCDF4.Variable.__getitem__

check_var_keys.py

from dataclasses import dataclass
from typing import Any

import netCDF4
import numpy as np

# data
arr1_shape = (3,)
arr2_shape = (3, 10)
arr1 = np.arange(3, dtype=np.int64) + 5
arr2 = np.arange(30, dtype=np.int64).reshape(arr2_shape) + 10


@dataclass
class KeyTestCase:
    varname: str
    name: str
    nc_key: Any
    getitem_result: Any = None

    @staticmethod
    def csv_header():
        return '"key type", "key", "`__getitem__` result", "Comment"'

    @staticmethod
    def md_table_header():
        return (
            "| key type | key | `ds[varname][key]` result | Comment | \n"
            "|----------|-----|---------------------------|---------|"
        )

    def str_parts(self):
        # nc_key_str = f"`{type(self.nc_key).__name__}({self.nc_key})`"
        nc_key_str = f"`{repr(self.nc_key)}`"
        result_repr = repr(self.getitem_result)
        if isinstance(self.getitem_result, Exception):
            getitem_result = " ".join(result_repr.split())
        else:
            getitem_result = "<br />".join(f"`{line}`" for line in result_repr.split("\n"))
        comment = ""  # to be filled in by hand later
        return self.name, nc_key_str, getitem_result, comment

    def csv_line(self):
        return f'"{'","'.join(self.str_parts())}"'

    def md_table_line(self):
        return f"| {' | '.join(self.str_parts())} |"


var1_tests: list[KeyTestCase] = [
    KeyTestCase("var1", "int", 0),
    KeyTestCase("var1", "integral float", 0.0),
    KeyTestCase("var1", "Non-integral float", 0.8),
    KeyTestCase("var1", "int str", "1"),
    KeyTestCase("var1", "integral float string", "1.0"),
    KeyTestCase("var1", "Non-integral float string", "1.5"),
    KeyTestCase("var1", "bool", True),
    KeyTestCase("var1", "list of int", [0, 1]),
    KeyTestCase("var1", "list of bool (same shape as dimension)", [True, False, True]),
    KeyTestCase("var1", "list of bool (wrong shape)", [True, False]),
    KeyTestCase("var1", "list of integral float", [0.0, 1.0]),
    KeyTestCase("var1", "list of non-integral float", [0.6, 1.7]),
    KeyTestCase("var1", "list of int str", ["0", "1"]),
    KeyTestCase("var1", "list of integral float str", ["0.0", "1.0"]),
    KeyTestCase("var1", "list of non-integral float str", ["0.6", "1.7"]),
    KeyTestCase("var1", "single-valued list of int", [1]),
    KeyTestCase("var1", "single-valued list of integral float", [1.0]),
    KeyTestCase("var1", "single-valued list of non-integral float", [1.7]),
    KeyTestCase("var1", "single-valued list of int str", ["1"]),
    KeyTestCase("var1", "single-valued list of integral float str", ["1.0"]),
    KeyTestCase("var1", "single-valued list of non-integral float str", ["1.7"]),
]

var2_tests = [
    KeyTestCase("var2", "int", 0),
    KeyTestCase("var2", "integral float", 0.0),
    KeyTestCase("var2", "Non-integral float", 0.8),
    KeyTestCase("var2", "integer string", "1"),
    KeyTestCase("var2", "tuple of slice and int", (slice(None), 1)),
    KeyTestCase("var2", "tuple of ints", (0, 1)),
    KeyTestCase("var2", "tuple of integral floats", (0.0, 1.0)),
    KeyTestCase("var2", "tuple of non-integral floats", (0.8, 1.7)),
    KeyTestCase("var2", "tuple of lists of ints", ([0, 1], [1, 2])),
    KeyTestCase("var2", "tuple of lists of integral floats", ([0.0, 1.0], [1.0, 2.0])),
    KeyTestCase("var2", "tuple of lists of non-integral floats", ([0.8, 1.7], [1.4, 2.1])),
    KeyTestCase("var2", "tuple of lists of integer strings", (["0", "1"], ["1", "2"])),
    KeyTestCase("var2", "tuple of lists of integral float strings", (["0.0", "1.0"], ["1.0", "2.0"])),
    KeyTestCase("var2", "tuple of lists of non-integral float strings", (["0.8", "1.7"], ["1.4", "2.1"])),
    KeyTestCase("var2", "tuple of slice and integral float", (slice(None), 1.0)),
    KeyTestCase("var2", "tuple of slice and non-integral float", (slice(None), 1.8)),
    KeyTestCase("var2", "tuple of slice and integer string", (slice(None), "1")),
    KeyTestCase("var2", "tuple of slice and integral float string", (slice(None), "1.0")),
    KeyTestCase("var2", "tuple of slice and non-integral float string", (slice(None), "1.8")),
    KeyTestCase("var2", "tuple of slice and bool", (slice(None), True)),
    KeyTestCase("var2", "tuple of slice and list of int", (slice(None), [1, 2])),
    KeyTestCase("var2", "tuple of slice and list of integral float", (slice(None), [1.0, 2.0])),
    KeyTestCase("var2", "tuple of slice and list of bool (same shape as dimension)", ([True, False, True], slice(None))),
]

with netCDF4.Dataset("test.nc", "w") as ds:
    dim1 = ds.createDimension("d1", arr2_shape[0])
    dim2 = ds.createDimension("d2", arr2_shape[1])
    var1 = ds.createVariable("var1", np.int64, dim1)
    var1.set_always_mask(False)
    var2 = ds.createVariable("var2", np.int64, (dim1, dim2))
    var2.set_always_mask(False)
    var1[...] = arr1
    var2[...] = arr2
    var1_repr = repr(var1[...])
    var2_repr = repr(var2[...])
    for test in [*var1_tests, *var2_tests]:
        try:
            test.getitem_result = ds[test.varname][test.nc_key]
        except Exception as exc:
            test.getitem_result = exc

_1d_header = f"""
## 1-D variable tests

`arr1`:

```python
{repr(arr1)}
```

`var1[...]`:

```python,
{var1_repr}
```
"""

_2d_header = f"""
## 2-D variable tests

`arr2`:

```python
{repr(arr2)}
```

`var1[...]`:

```python,
{var2_repr}
```
"""

print("\n".join([
    _1d_header,
    "",
    KeyTestCase.md_table_header(),
    *[test.md_table_line() for test in var1_tests],
    "",
    _2d_header,
    "",
    KeyTestCase.md_table_header(),
    *[test.md_table_line() for test in var2_tests],
]))