fix include_rest support (elastic#20)

pheyos · May 11, 2020 · 5896b89 · 5896b89
1 parent 69f5ccc
commit 5896b89
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 6 deletions.
diff --git a/anonymize_it/anonymize.py b/anonymize_it/anonymize.py
@@ -38,4 +38,4 @@ def read_config(config_file):
     anon = Anonymizer(reader=reader, writer=writer)
 
     logging.info("performing anonymization...")
-    anon.anonymize(infer=True)
+    anon.anonymize(infer=True, include_rest=config.include_rest)
diff --git a/anonymize_it/anonymizers.py b/anonymize_it/anonymizers.py
@@ -49,7 +49,6 @@ def __init__(self, reader=None, writer=None, field_maps={}):
         self.suppressed_fields = None
         self.reader_type = None
         self.writer_type = None
-        self.include_rest = None
 
         # only parse config if it exists
         # this allows for anonymizer class instantiation via direct parameter setting
@@ -84,7 +83,7 @@ def instantiate_writer(self):
 
         self.writer = writer(dest_params)
 
-    def anonymize(self, infer=False):
+    def anonymize(self, infer=False, include_rest=False):
         """this is the core method for anonymizing data
 
         it utilizes specific reader and writer class methods to retrieve and store data. in the process
@@ -109,7 +108,7 @@ def anonymize(self, infer=False):
         total = self.reader.get_count()
         logging.info("total number of records {}...".format(total))
 
-        data = self.reader.get_data(list(self.field_maps.keys()), self.suppressed_fields, self.include_rest)
+        data = self.reader.get_data(list(self.field_maps.keys()), self.suppressed_fields, include_rest)
 
         # batch process the data and write out to json in chunks
         count = 0

diff --git a/anonymize_it/utils.py b/anonymize_it/utils.py
@@ -46,8 +46,8 @@ def parse_config(config):
     if not writer_type:
         raise ConfigParserError("destination error: dest type not defined. Please check config.")
 
-    Config = collections.namedtuple('Config', 'source dest masked_fields suppressed_fields')
-    config = Config(source, dest, masked_fields, suppressed_fields)
+    Config = collections.namedtuple('Config', 'source dest masked_fields suppressed_fields include_rest')
+    config = Config(source, dest, masked_fields, suppressed_fields, include_rest)
     return config