Units are now reported properly on NDFrame operations if all columns …

…have the same units (#26)

* Reports units for frame ops if units are all the same for each column

* isort black

energy_pandas/__init__.py

@@ -1,9 +1,9 @@
 """energy-pandas Module."""
 
 from outdated import warn_if_outdated
+from pkg_resources import DistributionNotFound, get_distribution
 
-from .energypandas import EnergySeries, EnergyDataFrame
-from pkg_resources import get_distribution, DistributionNotFound
+from .energypandas import EnergyDataFrame, EnergySeries
 
 # Version of the package
 try:

energy_pandas/energypandas.py

@@ -32,7 +32,7 @@
     _setup_subplots,
     save_and_show,
 )
-from .units import unit_registry, IP_DEFAULT_CONVERSION, SI_DEFAULT_CONVERSION
+from .units import IP_DEFAULT_CONVERSION, SI_DEFAULT_CONVERSION, unit_registry
 
 log = logging.getLogger(__name__)
 
@@ -143,7 +143,11 @@ def __finalize__(self, other, method=None, **kwargs):
             for name in other._metadata:
                 if name == "units":
                     if isinstance(other, EnergyDataFrame):
-                        setattr(self, name, getattr(other, "units").get(self.name))
+                        if len(set(getattr(other, "units").values())) == 1:
+                            # if all the same units
+                            setattr(self, name, *set(getattr(other, "units").values()))
+                        else:
+                            setattr(self, name, getattr(other, "units").get(self.name))
                     else:
                         setattr(self, name, getattr(other, "units"))
                 elif name == "name":

energy_pandas/plotting.py

@@ -1,5 +1,6 @@
 """Plotting module."""
 
+import logging
 import os
 import time
 
@@ -10,8 +11,6 @@
 from numpy.core._asarray import asarray
 from pandas.plotting._matplotlib.tools import create_subplots, flatten_axes
 
-import logging
-
 logger = logging.getLogger(__name__)
 log = logger.info
 

energy_pandas/units.py

@@ -4,9 +4,9 @@
 from tokenize import TokenInfo
 
 import pint
+import pkg_resources
 from pint.compat import tokenizer
 
-import pkg_resources
 energyplus_registry = pkg_resources.resource_filename(__name__, "energyplus_en.txt")
 
 

tests/test_energypandas.py

@@ -3,7 +3,7 @@
 from pandas import date_range
 
 from energy_pandas import EnergyDataFrame, EnergySeries
-from energy_pandas.units import unit_registry, dash_to_mul, underline_dash
+from energy_pandas.units import dash_to_mul, underline_dash, unit_registry
 
 
 @pytest.fixture()
@@ -66,7 +66,7 @@ def test_units(self, es):
         assert es.units == unit_registry.degF
         assert type(es) == EnergySeries
 
-        #set attribute
+        # set attribute
         es.units = "degC"
         assert es.units == unit_registry.degC
 
@@ -192,10 +192,16 @@ def test_mixed_units(self, edf_from_e_series):
             "Series 1 degC": unit_registry.degC
         }
 
+    def test_mixed_units_ops(self, edf_from_e_series):
+        col_1 = edf_from_e_series.iloc[:, 0]
+        col_2 = edf_from_e_series.iloc[:, 1]
+        # todo: deal with mixed units magnitude
+        assert (col_1 * col_2).units == "degree_Celsius"
+
     def test_mixed_units_convert(self, edf_from_e_series):
         assert edf_from_e_series.to_units("degR").units == {
             "Series 1 degC": unit_registry.degR,
-            "Series 2 degK": unit_registry.degR
+            "Series 2 degK": unit_registry.degR,
         }
 
     def test_units_value(self, edf):
@@ -235,6 +241,10 @@ def test_slice(self, edf):
         # check that the slice keeps the units
         assert edf.units == {"Temp": edf["Temp"].units}
 
+    def test_numeric_operations(self, edf):
+        assert edf.mean(axis=1).units == "degree_Celsius"
+        assert edf.sum(axis=1).units == "degree_Celsius"
+
     def test_repr(self, edf):
         # check that a slice returns an EnergySeries
         print(edf.__repr__())