Repository: wtbarnes/fiasco

Issue:
Allow slicing when iterating through levels on Ion
Currently this works,

```
ion = fiasco.Ion('H 0+', 1*u.MK)
ion[0]
```

but this does not

```
ion[:5]
```

We need to sort out some way to get slicing to work so that the second statement returns a list of the first four levels.

Discussion:

Patch:
diff --git a/fiasco/ion.py b/fiasco/ion.py
index c118d1a..ab77c6c 100644
--- a/fiasco/ion.py
+++ b/fiasco/ion.py
@@ -95,15 +95,27 @@ Using Datasets:
   abundance: {self._dset_names['abundance_filename']}
   ip: {self._dset_names['ip_filename']}"""
 
-    def __getitem__(self, key):
+    @cached_property
+    def _all_levels(self):
         try:
             _ = self._elvlc
         except KeyError:
+            return None
+        else:
+            n_levels = self._elvlc['level'].shape[0]
+            return [Level(i, self._elvlc) for i in range(n_levels)]
+
+    def __getitem__(self, key):
+        if self._all_levels is None:
             raise IndexError(f'No energy levels available for {self.ion_name}')
         else:
-            # Throw an index error to stop iteration
-            _ = self._elvlc['level'][key]
-            return Level(key, self._elvlc)
+            # NOTE: casting to array first so that "fancy indexing" can
+            # be used to index the energy levels.
+            indexed_levels = np.array(self._all_levels)[key]
+            if isinstance(indexed_levels, Level):
+                return indexed_levels
+            else:
+                return indexed_levels.tolist()
 
     def __add__(self, value):
         return IonCollection(self, value)


Test patch:
diff --git a/fiasco/tests/test_ion.py b/fiasco/tests/test_ion.py
index 5066511..a5c240b 100644
--- a/fiasco/tests/test_ion.py
+++ b/fiasco/tests/test_ion.py
@@ -6,6 +6,7 @@ import astropy.units as u
 import pytest
 
 import fiasco
+from fiasco.level import Level
 from fiasco.util.exceptions import MissingDatasetException
 
 temperature = np.logspace(5, 8, 100)*u.K
@@ -39,6 +40,24 @@ def test_new_instance(ion):
     assert ion._instance_kwargs['abundance_filename'] == abundance_filename
 
 
+def test_level_indexing(ion):
+    # Integer
+    assert isinstance(ion[0], Level)
+    assert ion[0].__repr__() == Level(0, ion._elvlc).__repr__()
+    # Slice
+    levels = ion[:5]
+    assert len(levels) == 5
+    assert isinstance(levels, list)
+    assert isinstance(levels[0], Level)
+    assert levels[2].__repr__() == Level(2, ion._elvlc).__repr__()
+    # Fancy indexing
+    levels = ion[[1, 5, 10]]
+    assert len(levels) == 3
+    assert isinstance(levels, list)
+    assert isinstance(levels[0], Level)
+    assert levels[2].__repr__() == Level(10, ion._elvlc).__repr__()
+
+
 def test_level(ion):
     level = ion[0]
     assert isinstance(level, fiasco.ion.Level)