Skip to content

Commit

Permalink
Implemented dict-like access for Stokes axes
Browse files Browse the repository at this point in the history
As per the second suggestion in radio-astro-tools#750 - the Stokes data can now be
accessed via cube['I'], cube['Q'] etc.

Assigning new keys is not yet supported and will throw
NotImplementedError if attempted.
  • Loading branch information
Kitchi committed Nov 2, 2021
1 parent feecc04 commit 0a5f2d4
Showing 1 changed file with 13 additions and 0 deletions.
13 changes: 13 additions & 0 deletions spectral_cube/stokes_spectral_cube.py
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,19 @@ def __init__(self, stokes_data, mask=None, meta=None, fill_value=None):

self._mask = mask

def __getitem__(self, key):
if key in self._stokes_data:
return self._stokes_data[key]
else:
raise KeyError("Key {0} does not exist in this cube.".format(key))

def __setitem__(self, key, item):
if key in self._stokes_data:
self._stokes_data[key] = item
else:
errmsg = "Assigning new Stokes axes is not yet supported."
raise NotImplementedError(errmsg)

@property
def shape(self):
return self._shape
Expand Down

0 comments on commit 0a5f2d4

Please sign in to comment.