Skip to content

Commit

Permalink
Add getitem support.
Browse files Browse the repository at this point in the history
  • Loading branch information
riga committed Aug 7, 2024
1 parent f5e2eeb commit b634b2a
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 1 deletion.
15 changes: 14 additions & 1 deletion scinum/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -840,7 +840,7 @@ def copy(
if uncertainties is None:
uncertainties = self.uncertainties

return self.__class__(nominal, uncertainties=uncertainties)
return self.__class__(nominal, uncertainties=uncertainties, **self._init_kwargs()) # type: ignore[arg-type] # noqa

def get(
self,
Expand Down Expand Up @@ -1113,6 +1113,19 @@ def __call__(
# shorthand for get
return self.get(direction=direction, names=names, unc=unc, factor=factor)

def __getitem__(self, key: int | slice) -> Number:
if not self.is_numpy:
raise TypeError("cannot index non-NumPy number")

return self.__class__(
self.nominal[key],
uncertainties={
name: (up[key], down[key])
for name, (up, down) in self.uncertainties.items()
},
**self._init_kwargs(), # type: ignore[arg-type]
)

def __float__(self) -> OutValueType:
# extract nominal value
return self.nominal
Expand Down
8 changes: 8 additions & 0 deletions tests/test_number.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,14 @@ def test_copy_numpy(self) -> None:
self.assertEqual(len(num.uncertainties), 1)
self.assertEqual(num.u(direction=UP).shape, (3,))

@if_numpy
def test_getitem_numpy(self) -> None:
num = Number(np.array([5, 27, 42]), 5)[1:]

self.assertEqual(tuple(num.nominal), (27, 42))
self.assertEqual(len(num.uncertainties), 1)
self.assertEqual(tuple(num.uncertainties["default"][0]), (5, 5))

def test_string_formats(self) -> None:
self.assertEqual(len(self.num.str()), 102)
self.assertEqual(len(self.num.str("%.3f")), 126)
Expand Down

0 comments on commit b634b2a

Please sign in to comment.