Skip to content

Commit

Permalink
fix issue
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewgsavage committed Jul 29, 2024
1 parent 1d1d40a commit ed2b198
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 12 deletions.
21 changes: 10 additions & 11 deletions pint_pandas/pint_array.py
Original file line number Diff line number Diff line change
Expand Up @@ -214,19 +214,18 @@ def _get_common_dtype(self, dtypes):
-------
returns self for acceptable cases or None otherwise
"""
# Return self (PintType with same units) if possible
if all(
isinstance(x, PintType) or pd.api.types.is_numeric_dtype(x) for x in dtypes
isinstance(dtype, PintType) and dtype.units == self.units
for dtype in dtypes
):
PintType_list = [x for x in dtypes if isinstance(x, PintType)]
if len(PintType_list) < 2:
return self
if all(
PintType_list[0].units.is_compatible_with(x.units)
for x in PintType_list[1:]
):
return self
else:
return None
return self
# Otherwise return PintType with undefined units
elif all(
isinstance(dtype, PintType) or pd.api.types.is_numeric_dtype(dtype)
for dtype in dtypes
):
return PintType
else:
return None

Expand Down
25 changes: 24 additions & 1 deletion pint_pandas/testsuite/test_issues.py
Original file line number Diff line number Diff line change
Expand Up @@ -277,6 +277,7 @@ def test_eval(self):
)
tm.assert_series_equal(df.eval("a / b"), df["a"] / df["b"])
tm.assert_series_equal(df.eval("a / c"), df["a"] / df["c"])
tm.assert_series_equal(df.eval("a / a"), df["a"] / df["a"])

def test_mixed_df(self):
df = pd.DataFrame(
Expand All @@ -287,4 +288,26 @@ def test_mixed_df(self):
}
)

assert df["a"][0] == df.iloc[0][0]
assert df["a"][0] == df.iloc[0, 0]


class TestIssue246(BaseExtensionTests):
def test_issue246(self):
df = pd.DataFrame(
{
"a": [1, 2, 3],
"b": [4, 5, 6],
"c": [7, 8, 9],
}
)

df = df.astype(
{
"a": "pint[m]",
"b": "pint[m/s]",
"c": "pint[kN]",
}
)

# now an operation where each cell is independent from each other
df.apply(lambda x: x * 2, axis=1)

0 comments on commit ed2b198

Please sign in to comment.