Skip to content

Commit

Permalink
Make filter work for np 64 format years/dates (#370)
Browse files Browse the repository at this point in the history
* Tests and code added

* Added documentation

* Appease stickler

* Fixing for review comments

Co-authored-by: Lamboll <rlamboll@ic.ac.uk>
  • Loading branch information
Rlamboll and Lamboll authored Apr 21, 2020
1 parent 8ead13e commit b64263b
Show file tree
Hide file tree
Showing 4 changed files with 21 additions and 8 deletions.
2 changes: 1 addition & 1 deletion RELEASE_NOTES.md
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# Next Release

## Individual Updates

- [#370](https://github.com/IAMconsortium/pyam/pull/370) Allowed filter to work with np.int64 years and np.datetime64 dates.
- [#361](https://github.com/IAMconsortium/pyam/pull/361) iam-units refactored from a Git submodule to a Python dependency of pyam.

# Release v0.5.0
Expand Down
6 changes: 3 additions & 3 deletions pyam/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -1178,10 +1178,10 @@ def filter(self, keep=True, inplace=False, **kwargs):
string or list of strings, where `*` can be used as a wildcard
- 'level': the maximum "depth" of IAM variables (number of '|')
(excluding the strings given in the 'variable' argument)
- 'year': takes an integer, a list of integers or a range
note that the last year of a range is not included,
- 'year': takes an integer (int/np.int64), a list of integers or
a range. Note that the last year of a range is not included,
so `range(2010, 2015)` is interpreted as `[2010, ..., 2014]`
- arguments for filtering by `datetime.datetime`
- arguments for filtering by `datetime.datetime` or np.datetime64
('month', 'hour', 'time')
- 'regexp=True' disables pseudo-regexp syntax in `pattern_match()`
"""
Expand Down
12 changes: 8 additions & 4 deletions pyam/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -353,8 +353,10 @@ def _escape_regexp(s):

def years_match(data, years):
"""Return rows where data matches year"""
years = [years] if isinstance(years, int) else years
dt = datetime.datetime
years = [years] if (
isinstance(years, (int, np.int64))
) else years
dt = (datetime.datetime, np.datetime64)
if isinstance(years, dt) or isinstance(years[0], dt):
error_msg = "`year` can only be filtered with ints or lists of ints"
raise TypeError(error_msg)
Expand Down Expand Up @@ -423,9 +425,11 @@ def conv_strs(strs_to_convert, conv_codes, name):
def datetime_match(data, dts):
"""Matching of datetimes in time columns for data filtering"""
dts = dts if islistable(dts) else [dts]
if any([not isinstance(i, datetime.datetime) for i in dts]):
if any([not (
isinstance(i, (datetime.datetime, np.datetime64))
) for i in dts]):
error_msg = (
"`time` can only be filtered by datetimes"
"`time` can only be filtered by datetimes and datetime64s"
)
raise TypeError(error_msg)
return data.isin(dts)
Expand Down
9 changes: 9 additions & 0 deletions tests/test_core.py
Original file line number Diff line number Diff line change
Expand Up @@ -317,6 +317,15 @@ def test_filter_day(test_df, test_day):
assert unique_time[0] == expected


def test_filter_with_numpy_64_date_vals(test_df):
dates = test_df[test_df.time_col].unique()
key = 'year' if test_df.time_col == "year" else 'time'
res_0 = test_df.filter(**{key: dates[0]})
res = test_df.filter(**{key: dates})
assert np.equal(res_0.data[res_0.time_col].values, dates[0]).all()
assert res.equals(test_df)


@pytest.mark.parametrize("test_hour", [0, 12, [12, 13]])
def test_filter_hour(test_df, test_hour):
if "year" in test_df.data.columns:
Expand Down

0 comments on commit b64263b

Please sign in to comment.