fracfocus_filter_class.py

class fracfocus_data_search:
    """
    This class generates an object that is used the filter the master
    fracfocus dataframe so it only contains a certain state/state abbreviation,
    county (list), and/or operator
    """

    def __init__(self, state=None, state_abbreviation=None, county_list=None,
                 operator=None):
        #All data in initialize def optional depending on what kind of filtering
        #we want to do
        self.state = state
        self.state_abbreviation = state_abbreviation
        self.county_list=county_list
        self.operator=operator

    def filter_dataframe(self, df, column_state, 
                        column_county, column_operator):
        """
        Filter the data by parameters
        """
        #Filter by all the parameters and return the subset
        df_subset=df[
                #By state
                (((df[column_state].str.lower())==self.state.lower()) | 
                ((df[column_state].str.lower())==self.state_abbreviation.lower())) &
                #By county
                ((df[column_county].str.lower()).isin(map(str.lower,self.county_list))) &
                #By operator
                ((df[column_operator].str.lower()).str.contains(self.operator.lower()))  
                ]
        return df_subset