Skip to content

Instantly share code, notes, and snippets.

@janpipek
Last active January 26, 2020 02:59
Show Gist options
  • Save janpipek/d9e3cfe471300d40101ad61d744f3286 to your computer and use it in GitHub Desktop.
Save janpipek/d9e3cfe471300d40101ad61d744f3286 to your computer and use it in GitHub Desktop.
How to implement `which` in pandas
import pandas as pd
def which(series):
# Error handling omitted
if not isinstance(series, pd.Series):
series = pd.Series(series)
return series[series.astype(bool) == True].index.tolist()
# Way to extend a Series
# See: https://pandas.pydata.org/pandas-docs/stable/development/extending.html#registering-custom-accessors
@pd.api.extensions.register_series_accessor("r")
class RAcccessor:
def __init__(self, series):
self._series = series
def which(self):
return self._series[self._series.astype(bool) == True].index.tolist()
... # Add more accessors of your liking
pd.Series([True, False, True, False, False], index=["a", "c", "e", "g", "i"]).r.which()
@janpipek
Copy link
Author

janpipek commented Jan 23, 2020

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment