Skip to content

Instantly share code, notes, and snippets.

@RandallPittmanOrSt

RandallPittmanOrSt/check_var_keys.py

Last active July 19, 2024 22:38
Show Gist options
  • Save RandallPittmanOrSt/e0cc24becf19183cf293f515e583cf11 to your computer and use it in GitHub Desktop.
Save RandallPittmanOrSt/e0cc24becf19183cf293f515e583cf11 to your computer and use it in GitHub Desktop.
Download ZIP
Check various keys for netCDF4.Variable.__getitem__
Raw
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],
]))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment