refactor anonymizer with inheritance

anonymizer.py

@@ -1,67 +1,111 @@
-from typing import List
+import logging
+from pathlib import Path
+from typing import List, Optional
 
+import pandas as pd
 from presidio_analyzer import RecognizerResult
 from presidio_anonymizer import AnonymizerEngine
 from presidio_anonymizer.entities import OperatorConfig
+from presidio_anonymizer.operators import OperatorType
 
 from names_database import NameDatabase
 
-names_db = NameDatabase()
+name_table = Path('data', 'ascii_names.parquet')
 
-def split_name(original_name: str):
-    '''Splits name into parts.
-    If one token, assume it is a first name.
-    If two tokens, first and last name.
-    If three tokens, one first name and two last names.
-    If four tokens, two first names and two last names.'''
-    names = original_name.split()
-    if len(names) == 1:
-        return names[0], None
-    elif len(names) == 2:
-        return names[0], names[1]
-    elif len(names) == 3:
-        return names[0], ' '.join(names[1:])
-    elif len(names) == 4:
-        return ' '.join(names[:2]), ' '.join(names[2:])
-    else:
-        return None, None
+logger = logging.getLogger('anonymizer')
 
-def generate_surrogate(original_name: str):
-    '''Generate a surrogate name.
-    '''
-    first_names, last_names = split_name(original_name)
-    gender = names_db.get_gender(first_names) if first_names else None
-    country = names_db.get_country(last_names) if last_names else None
-    
-    surrogate_name = ''
-    
-    name_candidates = names_db.get_random_name(
-        gender=gender,
-        country=country)
-    
-    surrogate_name += name_candidates.iloc[0]['first']
-    
-    if last_names:
-        surrogate_name += ' ' + name_candidates.iloc[1]['last']
+
+class surrogate_anonymizer(AnonymizerEngine):
+    def __init__(self):
+        super().__init__()
+        self.names_db = NameDatabase()
+        self.names_df = pd.read_parquet(name_table)
+        
+        
+    def get_random_name(
+            self,
+            country: Optional[str] = None,
+            gender: Optional[str] = None
+    ) -> pd.DataFrame:
+        '''Returns two random names from the database as a DataFrame.
+        Both rows match gender and country, if provided.
+        :country: ISO country code e.g. "CO" for Columbia
+        :gender: 'M' or 'F'
+        returns two rows of the names dataframe
+        '''
+        names_view = self.names_df
+        if country:
+            names_view = names_view[names_view['country'] == country]
+        if gender:
+            names_view = names_view[names_view['gender'] == gender]
+        if names_view.size < 25:
+            return self.names_df.sample(n=2, weights=self.names_df['count'])
+        return names_view.sample(n=2, weights=names_view['count'])
+
+    def split_name(self, original_name: str):
+        '''Splits name into parts.
+        If one token, assume it is a first name.
+        If two tokens, first and last name.
+        If three tokens, one first name and two last names.
+        If four tokens, two first names and two last names.'''
+        names = original_name.split()
+        if len(names) == 1:
+            logger.info(f'Splitting to 1 first name: {names}')
+            return names[0], None
+        elif len(names) == 2:
+            logger.info(f'Splitting to 1 first name, 1 last name: {names}')
+            return names[0], names[1]
+        elif len(names) == 3:
+            logger.info(f'Splitting to 1 first name, 2 last names: {names}')
+            return names[0], ' '.join(names[1:])
+        elif len(names) == 4:
+            logger.info(f'Splitting to 2 first names and 2 last names: {names}')
+            return ' '.join(names[:2]), ' '.join(names[2:])
+        else:
+            logger.info(f'Splitting failed, do not match gender/country: {names}')
+            return None, None
+
+    def generate_surrogate(self, original_name: str):
+        '''Generate a surrogate name.
+        '''
+        first_names, last_names = self.split_name(original_name)
+        gender = self.names_db.get_gender(first_names) if first_names else None
+        logger.debug(f'Gender set to {gender}')
+        country = self.names_db.get_country(last_names) if last_names else None
+        logger.debug(f'Country set to {country}')
+        
+        surrogate_name = ''
         
-    return surrogate_name
+        name_candidates = self.get_random_name(gender=gender, country=country)
+        
+        surrogate_name += name_candidates.iloc[0]['first']
+        logger.info(f'First name surrogate is {surrogate_name}')
+        
+        if last_names:
+            logger.info(f'Combining with {name_candidates.iloc[1]["last"]}')
+            surrogate_name += ' ' + name_candidates.iloc[1]['last']
+            
+        logger.info(f'Returning surrogate name {surrogate_name}')
+        return surrogate_name
 
-def anonymize(
-    anonymizer: AnonymizerEngine,
-    text: str,
-    analyze_results: List[RecognizerResult]
-    ):
-    '''Anonymize identified input using Presidio Anonymizer.'''
-    
-    if not text:
-        return
-    
-    res = anonymizer.anonymize(
-        text,
-        analyze_results,
-        operators={
+    def anonymize(
+        self,
+        text: str,
+        analyzer_results: List[RecognizerResult]
+        ):
+        '''Anonymize identified input using Presidio Anonymizer.'''
+        
+        if not text:
+            return
+        
+        analyzer_results = self._remove_conflicts_and_get_text_manipulation_data(
+            analyzer_results
+        )
+        
+        operators = self._AnonymizerEngine__check_or_add_default_operator(
+            {
             'STUDENT': OperatorConfig('custom',
-                                      {'lambda': generate_surrogate}),
+                                      {'lambda': self.generate_surrogate}),
             'EMAIL_ADDRESS': OperatorConfig('replace',
                                             {'new_value': '[email protected]'}),
             'PHONE_NUMBER': OperatorConfig('replace',
@@ -69,9 +113,24 @@ def anonymize(
             'URL': OperatorConfig('replace',
                                   {'new_value': 'aol.com'}),
             }
-    )
-    
-    return res.text
+        )
+        
+        res = self._operate(text,
+                            analyzer_results,
+                            operators,
+                            OperatorType.Anonymize)
+                
+        return res.text
 
 if __name__ == '__main__':
-    print(generate_surrogate('Nora Wang'))
+    logging.basicConfig(level=logging.DEBUG)
+    anonymizer = surrogate_anonymizer()
+    test_names = ['Nora Wang',
+                  'MJ',
+                  '',
+                  '(',
+                  'Mario Escobar Sanchez',
+                  'Jane Fonda Michelle Rousseau',
+                  'Sir Phillipe Ricardo de la Sota Mayor']
+    for name in test_names:
+        anonymizer.generate_surrogate(name)

app.py

@@ -2,7 +2,7 @@
 '''Streamlit app for Student Name Detection models.'''
 
 from analyzer import prepare_analyzer
-from anonymizer import anonymize
+from anonymizer import surrogate_anonymizer
 from presidio_anonymizer import AnonymizerEngine
 import pandas as pd
 from annotated_text import annotated_text
@@ -31,8 +31,8 @@ def analyzer_engine():
 
 @st.cache(allow_output_mutation=True)
 def anonymizer_engine():
-    '''Return AnonymizerEngine.'''
-    return AnonymizerEngine()
+    '''Return generate surrogate anonymizer.'''
+    return surrogate_anonymizer()
 
 def annotate(text, st_analyze_results, st_entities):
     tokens = []
@@ -116,10 +116,9 @@ with st.spinner('Analyzing...'):
 st.text('')
 
 st.subheader('Anonymized')
-
 with st.spinner('Anonymizing...'):
     if button or st.session_state.first_load:
-        st_anonymize_results = anonymize(anonymizer_engine(),
+        st_anonymize_results = anonymizer_engine().anonymize(
                                          st_text,
                                          st_analyze_results)
         st_anonymize_results

names_database.py

@@ -1,42 +1,35 @@
-from pathlib import Path
-from typing import Optional
+import logging
+
 
-import pandas as pd
 from names_dataset import NameDataset, NameWrapper
 
-name_table = Path('data', 'ascii_names.parquet')
 
 class NameDatabase(NameDataset):
     def __init__(self) -> None:
         super().__init__()
-        self.names = pd.read_parquet(name_table)
-
-    def get_random_name(
-            self,
-            country: Optional[str] = None,
-            gender: Optional[str] = None
-    ):
-        '''country: ISO country code in 'alpha 2' format
-        gender: 'M' or 'F'
-        returns two rows of the names dataframe
-        '''
-        names_view = self.names
-        if country:
-            names_view = names_view[names_view['country'] == country]
-        if gender:
-            names_view = names_view[names_view['gender'] == gender]
-        if names_view.size < 25:
-            return self.names.sample(n=2, weights=self.names['count'])
-        return names_view.sample(n=2, weights=names_view['count'])
+        
+        self.logger = logging.getLogger('anonymizer')
     
-    def search(self, name: str):
+    def search(self, name: str) -> dict:
+        '''Returns all entries associated with a name string.
+        The name string can be multiple tokens. 
+        Both first and last names will be matched.
+        '''
         key = name.strip().title()
         fn = self.first_names.get(key) if self.first_names is not None else None
         ln = self.last_names.get(key) if self.last_names is not None else None
         return {'first_name': fn, 'last_name': ln}
        
-    def get_gender(self, first_names: str):
-        return NameWrapper(self.search(first_names)).gender
+    def get_gender(self, first_names: str) -> str:
+        '''Return the most frequent gender code for a specific last name,
+        or None if a match cannot be found.
+        '''
+        gender = NameWrapper(self.search(first_names)).gender
+        return gender if gender else None
 
-    def get_country(self, last_names: str):
-        return NameWrapper(self.search(last_names)).country
+    def get_country(self, last_names: str) -> str:
+        '''Return the most frequent country code for a specific last name,
+        or None if a match cannot be found.
+        '''
+        country = NameWrapper(self.search(last_names)).country
+        return country if country else None