diff --git a/saucebrush/filters.py b/saucebrush/filters.py index 7fee37a..380628f 100644 --- a/saucebrush/filters.py +++ b/saucebrush/filters.py @@ -215,6 +215,24 @@ class FieldModifier(FieldFilter): str(self._target_keys), str(self._filter_func)) +class FieldKeeper(Filter): + """ Filter that removes all but the given set of fields. + + FieldKeeper(('spam', 'eggs')) removes all bu tthe spam and eggs + fields from every record filtered. + """ + + def __init__(self, keys): + super(FieldKeeper, self).__init__() + self._target_keys = utils.str_or_list(keys) + + def process_record(self, record): + for key in record.keys(): + if key not in self._target_keys: + del record[key] + return record + + class FieldRemover(Filter): """ Filter that removes a given set of fields. diff --git a/saucebrush/tests/filters.py b/saucebrush/tests/filters.py index cbd9cd7..d0f3b85 100644 --- a/saucebrush/tests/filters.py +++ b/saucebrush/tests/filters.py @@ -3,7 +3,7 @@ import operator import types from saucebrush.filters import (Filter, YieldFilter, FieldFilter, SubrecordFilter, ConditionalPathFilter, - ConditionalFilter, FieldModifier, + ConditionalFilter, FieldModifier, FieldKeeper, FieldRemover, FieldMerger, FieldAdder, FieldCopier, FieldRenamer, Unique) @@ -200,6 +200,13 @@ class FilterTestCase(unittest.TestCase): {'a':2, 'b':10, 'c':200}] self.assert_filter_result(fm, expected_data) + def test_field_keeper(self): + fk = FieldKeeper(['c']) + + # check against expected results + expected_data = [{'c':3}, {'c':5}, {'c':100}] + self.assert_filter_result(fk, expected_data) + def test_field_remover(self): fr = FieldRemover(['a', 'b'])