Fix bug where KeySelector didn't pass along kwargs correctly

featureflow/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '2.2.1'
+__version__ = '2.2.3'
 
 from model import BaseModel, ModelExistsError
 

featureflow/extractor.py

@@ -194,8 +194,8 @@ class KeySelector(object):
     A mixin for Node-derived classes that allows the extractor to process a
     single key from the dictionary-like object it is passed
     """
-    def __init__(self, aspect_key, needs=None):
-        super(KeySelector, self).__init__(needs=needs)
+    def __init__(self, aspect_key, needs=None, **kwargs):
+        super(KeySelector, self).__init__(needs=needs, **kwargs)
         self.aspect_key = aspect_key
 
     def _enqueue(self, data, pusher):

featureflow/feature.py

@@ -54,11 +54,12 @@ def _fixup_needs(self):
                 self.needs[k] = v.feature
 
                 class AspectExtractor(KeySelector, self.extractor):
-                    pass
+                    def __init__(self, *args, **kwargs):
+                        print args, kwargs
+                        super(AspectExtractor, self).__init__(*args, **kwargs)
 
                 self.extractor = AspectExtractor
                 self.extractor_args.update(aspect_key=v.aspect_key)
-
             except AttributeError:
                 # the value is already just a feature
                 pass

featureflow/test_integration.py

@@ -215,17 +215,22 @@ def _process(self, data):
 
 
 class ExplicitConcatenate(Aggregator, Node):
-    def __init__(self, needs=None):
+    def __init__(self, blarg=None, needs=None):
         super(ExplicitConcatenate, self).__init__(needs=needs)
+        self.blarg = blarg
         self._cache = defaultdict(str)
 
     def _enqueue(self, data, pusher):
-        print 'ENQUEUE', data, pusher
         k = self._dependency_name(pusher)
         self._cache[k] += data
 
+    def _last_chunk(self):
+        if self.blarg:
+            yield 'BLARG!'
+        else:
+            yield ''
+
     def _process(self, data):
-        print 'PROCESS', data
         lhs = data['lhs']
         rhs = data['rhs']
         yield lhs + rhs
@@ -418,6 +423,23 @@ class Document(BaseModel, self.Settings):
 
         self.assertEqual('THIS IS A TEST.this is a test.', doc.cat.read())
 
+    def test_feature_with_aspect_dependencies_can_also_have_arguments(self):
+        class Document(BaseModel, self.Settings):
+            stream = Feature(TextStream, store=True)
+            both = Feature(UpperAndLower, needs=stream, store=False)
+            upper = Feature(ToUpper, needs=stream, store=False)
+            cat = Feature(
+                ExplicitConcatenate,
+                blarg=True,
+                needs=dict(lhs=upper, rhs=both.aspect('lower')),
+                store=False)
+
+        keyname = 'cased'
+        _id = Document.process(stream=keyname)
+        doc = Document(_id)
+
+        self.assertEqual('THIS IS A TEST.this is a test.BLARG!', doc.cat.read())
+
     def test_can_have_multiple_explicit_dependencies(self):
         class Split(BaseModel, self.Settings):
             stream = Feature(TextStream, store=False)