Repository: lbl-anp/becquerel

Issue:
replace calls to Spectrum() with self.__class__() etc.
Throughout `Spectrum` (and probably other classes) there are many points where a `Spectrum` method returns a new `Spectrum` object by doing `return Spectrum(...)`. However this is kind of bad practice as downstream users may derive their own classes from `Spectrum`, in which case `Spectrum` methods will return a `Spectrum` rather than the subclass.

The fix is simple: replace calls to `Spectrum()` (etc.) with `self.__class__()`

Discussion:

Patch:
diff --git a/becquerel/core/spectrum.py b/becquerel/core/spectrum.py
index feff92a..a343c78 100644
--- a/becquerel/core/spectrum.py
+++ b/becquerel/core/spectrum.py
@@ -811,9 +811,9 @@ class Spectrum:
             kwargs = {"cps": self.cps + other.cps}
 
         if self.is_calibrated and other.is_calibrated:
-            spect_obj = Spectrum(bin_edges_kev=self.bin_edges_kev, **kwargs)
+            spect_obj = self.__class__(bin_edges_kev=self.bin_edges_kev, **kwargs)
         else:
-            spect_obj = Spectrum(bin_edges_raw=self.bin_edges_raw, **kwargs)
+            spect_obj = self.__class__(bin_edges_raw=self.bin_edges_raw, **kwargs)
         return spect_obj
 
     def __sub__(self, other):
@@ -867,9 +867,9 @@ class Spectrum:
                 )
 
         if self.is_calibrated and other.is_calibrated:
-            spect_obj = Spectrum(bin_edges_kev=self.bin_edges_kev, **kwargs)
+            spect_obj = self.__class__(bin_edges_kev=self.bin_edges_kev, **kwargs)
         else:
-            spect_obj = Spectrum(bin_edges_raw=self.bin_edges_raw, **kwargs)
+            spect_obj = self.__class__(bin_edges_raw=self.bin_edges_raw, **kwargs)
         return spect_obj
 
     def _add_sub_error_checking(self, other):
@@ -990,9 +990,9 @@ class Spectrum:
             data_arg = {"cps": self.cps * multiplier}
 
         if self.is_calibrated:
-            spect_obj = Spectrum(bin_edges_kev=self.bin_edges_kev, **data_arg)
+            spect_obj = self.__class__(bin_edges_kev=self.bin_edges_kev, **data_arg)
         else:
-            spect_obj = Spectrum(bin_edges_raw=self.bin_edges_raw, **data_arg)
+            spect_obj = self.__class__(bin_edges_raw=self.bin_edges_raw, **data_arg)
         return spect_obj
 
     def attenuate(self, material, areal_density_gcm2: float, **kwargs):
@@ -1083,13 +1083,13 @@ class Spectrum:
         new_counts = np.random.binomial(old_counts, 1.0 / f)
 
         if self.is_calibrated:
-            return Spectrum(
+            return self.__class__(
                 counts=new_counts,
                 bin_edges_kev=self.bin_edges_kev,
                 livetime=new_livetime,
             )
         else:
-            return Spectrum(
+            return self.__class__(
                 counts=new_counts,
                 bin_edges_raw=self.bin_edges_raw,
                 livetime=new_livetime,
@@ -1296,7 +1296,7 @@ class Spectrum:
             "bin_edges_kev": combined_bin_edges,
             "livetime": self.livetime,
         }
-        obj = Spectrum(**kwargs)
+        obj = self.__class__(**kwargs)
         return obj
 
     def rebin(
@@ -1355,7 +1355,7 @@ class Spectrum:
             zero_pad_warnings=zero_pad_warnings,
             include_overflows=include_overflows,
         )
-        return Spectrum(
+        return self.__class__(
             counts=out_spec,
             uncs=np.sqrt(out_spec),
             bin_edges_kev=out_edges,


Test patch:
diff --git a/tests/spectrum_test.py b/tests/spectrum_test.py
index 595d6d8..ce96627 100644
--- a/tests/spectrum_test.py
+++ b/tests/spectrum_test.py
@@ -1242,3 +1242,28 @@ def test_spectrum_rebin_like():
     spec2_rebin = spec2.rebin_like(spec1)
     assert np.all(np.isclose(spec1.bin_edges_kev, spec2_rebin.bin_edges_kev))
     assert np.isclose(spec2.counts_vals.sum(), spec2_rebin.counts_vals.sum())
+
+
+# ----------------------------------------------
+#           Test Spectrum inheritance
+# ----------------------------------------------
+class DerivedSpectrum(bq.Spectrum):
+    pass
+
+
+derived_spec = DerivedSpectrum(make_data())
+
+
+@pytest.mark.parametrize(
+    "spec",
+    [
+        derived_spec + derived_spec,
+        derived_spec - derived_spec,
+        derived_spec * 2,
+        derived_spec / 2,
+        derived_spec.downsample(f=2),
+        derived_spec.combine_bins(f=2),
+    ],
+)
+def test_spectrum_inheritance(spec):
+    assert isinstance(spec, DerivedSpectrum)