fixed inheritance of Grid class (#56)

- use suggestion by @AstroMike in #56
- add test for using a DerivedClass

Author: orbeckst

gridData/core.py

@@ -230,7 +230,7 @@ def resample(self, edges):
         coordinates = ndmeshgrid(*midpoints)
         # feed a meshgrid to generate all points
         newgrid = self.interpolated(*coordinates)
-        return Grid(newgrid, edges)
+        return self.__class__(newgrid, edges)
 
     def resample_factor(self, factor):
         """Resample to a new regular grid.
@@ -611,76 +611,76 @@ def __ne__(self, other):
 
     def __add__(self, other):
         self.check_compatible(other)
-        return Grid(self.grid + _grid(other), edges=self.edges)
+        return self.__class__(self.grid + _grid(other), edges=self.edges)
 
     def __sub__(self, other):
         self.check_compatible(other)
-        return Grid(self.grid - _grid(other), edges=self.edges)
+        return self.__class__(self.grid - _grid(other), edges=self.edges)
 
     def __mul__(self, other):
         self.check_compatible(other)
-        return Grid(self.grid * _grid(other), edges=self.edges)
+        return self.__class__(self.grid * _grid(other), edges=self.edges)
 
     def __truediv__(self, other):
         # truediv will always do true division (in Python 2 and Python 3);
         # we use from __future__ include division everywhere
         self.check_compatible(other)
-        return Grid(self.grid / _grid(other), edges=self.edges)
+        return self.__class__(self.grid / _grid(other), edges=self.edges)
 
     def __div__(self, other):
         # in Python 2 only (without __future__.division): will do "classic division"
         # https://docs.python.org/2/reference/datamodel.html#object.__div__
         if not six.PY2:
             raise NotImplementedError("__div__ is only available in Python 2, use __truediv__")
         self.check_compatible(other)
-        return Grid(self.grid.__div__(_grid(other)), edges=self.edges)
+        return self.__class__(self.grid.__div__(_grid(other)), edges=self.edges)
 
     def __floordiv__(self, other):
         self.check_compatible(other)
-        return Grid(self.grid // _grid(other), edges=self.edges)
+        return self.__class__(self.grid // _grid(other), edges=self.edges)
 
     def __pow__(self, other):
         self.check_compatible(other)
-        return Grid(numpy.power(self.grid, _grid(other)), edges=self.edges)
+        return self.__class__(numpy.power(self.grid, _grid(other)), edges=self.edges)
 
     def __radd__(self, other):
         self.check_compatible(other)
-        return Grid(_grid(other) + self.grid, edges=self.edges)
+        return self.__class__(_grid(other) + self.grid, edges=self.edges)
 
     def __rsub__(self, other):
         self.check_compatible(other)
-        return Grid(_grid(other) - self.grid, edges=self.edges)
+        return self.__class__(_grid(other) - self.grid, edges=self.edges)
 
     def __rmul__(self, other):
         self.check_compatible(other)
-        return Grid(_grid(other) * self.grid, edges=self.edges)
+        return self.__class__(_grid(other) * self.grid, edges=self.edges)
 
     def __rtruediv__(self, other):
         self.check_compatible(other)
-        return Grid(_grid(other) / self.grid, edges=self.edges)
+        return self.__class__(_grid(other) / self.grid, edges=self.edges)
 
     def __rdiv__(self, other):
         # in Python 2 only (without __future__.division): will do "classic division"
         # https://docs.python.org/2/reference/datamodel.html#object.__div__
         if not six.PY2:
             raise NotImplementedError("__rdiv__ is only available in Python 2, use __rtruediv__")
         self.check_compatible(other)
-        return Grid(self.grid.__rdiv__(_grid(other)), edges=self.edges)
+        return self.__class__(self.grid.__rdiv__(_grid(other)), edges=self.edges)
 
     def __rfloordiv__(self, other):
         self.check_compatible(other)
-        return Grid(_grid(other) // self.grid, edges=self.edges)
+        return self.__class__(_grid(other) // self.grid, edges=self.edges)
 
     def __rpow__(self, other):
         self.check_compatible(other)
-        return Grid(numpy.power(_grid(other), self.grid), edges=self.edges)
+        return self.__class__(numpy.power(_grid(other), self.grid), edges=self.edges)
 
     def __repr__(self):
         try:
             bins = self.grid.shape
         except AttributeError:
             bins = "no"
-        return '<Grid with ' + str(bins) + ' bins>'
+        return '<{0} with {1!r} bins>'.format(self.__class__, bins)
 
 
 def ndmeshgrid(*arrs):

gridData/tests/test_grid.py

@@ -2,24 +2,24 @@
 import six
 
 import numpy as np
-from numpy.testing import assert_array_equal, assert_array_almost_equal
+from numpy.testing import (assert_array_equal, assert_array_almost_equal,
+                           assert_almost_equal)
 
 import pytest
 
 from gridData import Grid
 
-class TestGrid(object):
-    @staticmethod
-    @pytest.fixture(scope="class")
-    def data():
-        d = dict(
-            griddata=np.arange(1, 28).reshape(3, 3, 3),
-            origin=np.zeros(3),
-            delta=np.ones(3))
-        d['grid'] = Grid(d['griddata'], origin=d['origin'],
-                         delta=d['delta'])
-        return d
+@pytest.fixture(scope="class")
+def data():
+    d = dict(
+        griddata=np.arange(1, 28).reshape(3, 3, 3),
+        origin=np.zeros(3),
+        delta=np.ones(3))
+    d['grid'] = Grid(d['griddata'], origin=d['origin'],
+                     delta=d['delta'])
+    return d
 
+class TestGrid(object):
     def test_init(self, data):
         g = Grid(data['griddata'], origin=data['origin'],
                  delta=1)
@@ -47,7 +47,7 @@ def test_addition(self, data):
         g = g + data['grid']
         assert_array_equal(g.grid.flat, (2 + (2 * data['griddata'])).flat)
 
-    def test_substraction(self, data):
+    def test_subtraction(self, data):
         g = data['grid'] - data['grid']
         assert_array_equal(g.grid.flat, np.zeros(27))
         g = 2 - data['grid']
@@ -140,3 +140,19 @@ def test_resample_factor(self, data):
         # check that the edges are the same
         assert_array_almost_equal(g.grid[::5, ::5, ::5],
                                   data['grid'].grid[::2, ::2, ::2])
+
+
+def test_inheritance(data):
+    class DerivedGrid(Grid):
+        pass
+
+    dg = DerivedGrid(data['griddata'], origin=data['origin'],
+                     delta=data['delta'])
+    result = dg + dg - 2.5 * dg / (dg + 5.3)
+
+    assert isinstance(result, DerivedGrid)
+
+    g = data['grid']
+    ref = g + g - 2.5 * g / (dg + 5.3)
+    assert_almost_equal(result.grid, ref.grid)
+