class Var(object):
    def __init__(self, name):
        super(Var, self).__init__()
        self.name = name class DataWriter(BaseDataWriter):
    def __init__(
            self,
            needs=None,
            _id=None,
            feature_name=None,
            feature_version=None,
            key_builder=None,
            database=None,
            event_log=None):
        super(DataWriter, self).__init__(
            needs=needs,
            key_builder=key_builder,
            database=database)

        self.event_log = event_log
        self.feature_version = feature_version
        self._id = _id
        self.feature_name = feature_name
        self.content_type = needs.content_type

    @property
    def key(self):
        assert self.feature_version is not None
        return self.key_builder.build(
            self._id, self.feature_name, self.feature_version)

    def __enter__(self):
        self._stream = self.database.write_stream(self.key, self.content_type)
        return self

    def _close_stream(self):
        try:
            self._stream.close()
        except (IOError, ValueError):
            pass

    def _cleanup_after_error(self):
        try:
            del self.database[self.key]
        except:
            pass

    def _log_events(self):
        if self.event_log is None:
            return

        self.event_log.append(
            json.dumps({
                '_id': self._id,
                'name': self.feature_name,
                'version': self.feature_version
            }))

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._close_stream()
        if exc_type:
            self._cleanup_after_error()
            return
        self._log_events()

    def _process(self, data):
        try:
            data = data.encode()
        except AttributeError:
            pass
        yield self._stream.write(data) class BytesIODataWriter(BaseDataWriter):
    def __init__(
            self,
            needs=None,
            _id=None,
            feature_name=None,
            feature_version=None,
            key_builder=None,
            database=None,
            event_log=None,
            buffer_size_limit=None):
        
        super(BytesIODataWriter, self).__init__(
            needs=needs,
            key_builder=key_builder,
            database=database)

        # set the buffer size limit to int32 max size if it isn't explicitly
        # provided
        self.buffer_size_limit = buffer_size_limit or (2 ** 31)
        self.event_log = event_log
        self.feature_version = feature_version
        self._id = _id
        self.feature_name = feature_name
        self.content_type = needs.content_type
        self._stream = BytesIO()

    def __exit__(self, exc_type, exc_val, exc_tb):
        del self._stream

    def _process(self, data):
        if len(data) < self.buffer_size_limit:
            try:
                yield self._stream.write(data)
            except TypeError:
                yield self._stream.write(data.encode())
        else:
            chunksize = self.buffer_size_limit // 2
            for i in range(0, len(data), chunksize):
                yield self._stream.write(data[i: i + chunksize]) class ClobberDataWriter(DataWriter):
    def __init__(
            self,
            needs=None,
            _id=None,
            feature_name=None,
            feature_version=None,
            key_builder=None,
            database=None,
            event_log=None):
        super(ClobberDataWriter, self).__init__(
            needs=needs,
            _id=_id,
            feature_name=feature_name,
            feature_version=feature_version,
            key_builder=key_builder,
            database=database,
            event_log=event_log)

    def _cleanup_after_error(self):
        # since we're overwriting previous iterations of the data on each
        # chunk, we expect the existing version to be valid, so we're not
        # going to delete anything
        pass

    def _process(self, data):
        self._stream = self.database.write_stream(self.key, self.content_type)
        try:
            x = self._stream.write(data)
        except TypeError:
            x = self._stream.write(data.encode())
        self._stream.close()
        yield x class JSONDecoder(GreedyTextDecoder):
    """
    A decoder that interprets the data as JSON
    """

    def __init__(self):
        super(JSONDecoder, self).__init__()

    def __call__(self, flo):
        s = super(JSONDecoder, self).__call__(flo)
        return json.loads(s)

    def __iter__(self, flo):
        yield self(flo) class Decoder(object):
    """
    The simplest possible decoder takes a file-like object and returns it
    """

    def __init__(self):
        super(Decoder, self).__init__()

    def __call__(self, flo):
        return flo

    def __iter__(self, flo):
        for chunk in chunked(flo):
            yield chunk class GreedyTextDecoder(TextDecoder):
    def __init__(self):
        super().__init__()

    def __call__(self, flo):
        return flo.read().decode()

    def __iter__(self, flo):
        yield self(flo) class DecoderNode(Node):
    def __init__(self, needs=None, decodifier=None, version=None):
        super(DecoderNode, self).__init__(needs=needs)
        self._version = version
        self.decoder = decodifier

    @property
    def version(self):
        return self._version

    def _process(self, data):
        for x in self.decoder.__iter__(data):
            yield x class BZ2Decoder(Decoder):
    """
    A decoder that decompresses data using the bz2 compression algorithm
    """

    def __init__(self):
        super(BZ2Decoder, self).__init__()

    def __call__(self, flo):
        return self.__iter__(flo)

    def __iter__(self, flo):
        decompressor = bz2.BZ2Decompressor()
        for chunk in chunked(flo):
            yield decompressor.decompress(chunk) class PickleDecoder(GreedyDecoder):
    def __init__(self):
        super(PickleDecoder, self).__init__()

    def __call__(self, flo):
        return loads(flo.read())

    def __iter__(self, flo):
        yield self(flo) class IdentityEncoder(Node):
    content_type = 'application/octet-stream'

    def __init__(self, needs=None):
        super(IdentityEncoder, self).__init__(needs=needs)

    def _seed(self):
        return b''

    def _enqueue(self, data, pusher):
        self._cache = data if data else self._seed() class JSONEncoder(Node):
    content_type = 'application/json'

    def __init__(self, needs=None):
        super(JSONEncoder, self).__init__(needs=needs)

    def _process(self, data):
        yield json.dumps(data) class TextEncoder(IdentityEncoder):
    content_type = 'text/plain'

    def __init__(self, needs=None):
        super(TextEncoder, self).__init__(needs=needs)

    def _seed(self):
        return ''

    def _process(self, data):
        yield data.encode() class BZ2Encoder(Node):
    content_type = 'application/octet-stream'

    def __init__(self, needs=None):
        super(BZ2Encoder, self).__init__(needs=needs)
        self._compressor = None
        self._cache = b''

    def _finalize(self, pusher):
        self._cache = b''

    def _first_chunk(self, data):
        self._compressor = bz2.BZ2Compressor()
        return data

    def _last_chunk(self):
        yield self._compressor.flush()

    def _process(self, data):
        compressed = self._compressor.compress(data)
        if compressed:
            yield compressed class PickleEncoder(Node):
    content_type = 'application/octet-stream'

    def __init__(self, needs=None):
        super(PickleEncoder, self).__init__(needs=needs)

    def _process(self, data):
        yield dumps(data, protocol=HIGHEST_PROTOCOL, byref=True) class Graph(dict):
    def __init__(self, **kwargs):
        super(Graph, self).__init__(**kwargs)

    def roots(self):
        return dict((k, v) for k, v in list(self.items()) if v.is_root)

    def leaves(self):
        return dict((k, v) for k, v in list(self.items()) if v.is_leaf)

    def subscriptions(self):
        subscriptions = defaultdict(list)
        for node in list(self.values()):
            for n in node.dependencies:
                subscriptions[id(n)].append(node)
        return subscriptions

    def remove_dead_nodes(self, features):
        # starting from the leaves, remove any nodes that are not stored, and 
        # have no stored consuming nodes
        mapping = dict((self[f.key], f) for f in features)
        nodes = deque(list(self.leaves().values()))
        while nodes:
            extractor = nodes.pop()
            nodes.extendleft(extractor.dependencies)
            try:
                feature = mapping[extractor]
            except KeyError:
                continue
            if extractor.is_leaf and not feature.store:
                extractor.disconnect()
                try:
                    del self[feature.key]
                except KeyError:
                    pass

    def process(self, **kwargs):
        # get all root nodes (those that produce data, rather than consuming 
        # it)
        roots = self.roots()

        # ensure that kwargs contains *at least* the arguments needed for the
        # root nodes
        intersection = set(kwargs.keys()) & set(roots.keys())
        if len(intersection) < len(roots):
            raise KeyError(
                (
                    'the keys {kw} were provided to the process() method, but the' \
                    + ' keys for the root extractors were {r}') \
                    .format(kw=list(kwargs.keys()), r=list(roots.keys())))

        graph_args = dict((k, kwargs[k]) for k in intersection)

        subscriptions = self.subscriptions()
        queue = deque()

        # get a generator for each root node.
        with ExitStack() as stack:
            [stack.enter_context(v) for v in self.values()]
            generators = [roots[k].process(v, queue=queue)
                          for k, v in list(graph_args.items())]
            for _ in zip_longest(*generators):
                while queue:
                    key, fname, kwargs = queue.pop()
                    for subscriber in subscriptions[key]:
                        func = getattr(subscriber, fname)
                        result = func(**kwargs)
                        if result is None:
                            continue
                        [_ for _ in result] class FunctionalNode(Node):
    def __init__(self, func, needs=None, closure_fingerprint=None):
        """
        `FunctionalNode` makes it possible to specify a stateless processing
        node using a simple Python function (or callable).

        Args:
            func (callable): A callable to stand in for the `_process()` method
            needs (Node or List of Node): Processing nodes on which this one
                depends
            closure_fingerprint (callable): An optional callable that takes a
                dictionary of closed-over variable names and their values, and
                returns a string or buffer object that may be added to the hash
                that will constitute the version number
        """
        super(FunctionalNode, self).__init__(needs=needs)
        self.closure_fingerprint = closure_fingerprint
        self.func = func

    @property
    def version(self):
        """
        Compute the version identifier for this functional node using the
        func code and local names.  Optionally, also allow closed-over variable
        values to affect the version number when closure_fingerprint is
        specified
        """
        try:
            f = self.func.__call__.__code__
        except AttributeError:
            f = self.func.__code__

        h = md5()
        h.update(f.co_code)
        h.update(str(f.co_names).encode())

        try:
            closure = self.func.__closure__
        except AttributeError:
            return h.hexdigest()

        if closure is None or self.closure_fingerprint is None:
            return h.hexdigest()

        d = dict(
            (name, cell.cell_contents)
            for name, cell in zip(f.co_freevars, closure))
        h.update(self.closure_fingerprint(d).encode())

        return h.hexdigest()

    def _process(self, data):
        result = self.func(data)
        if inspect.isgenerator(result):
            for x in result:
                yield x
        else:
            yield result class Node(object):
    def __init__(self, needs=None):
        super(Node, self).__init__()
        if not inspect.isgeneratorfunction(self._process):
            raise InvalidProcessMethod(self.__class__)

        self._cache = None
        self._listeners = []

        self._needs = dictify(needs)
        self._dependecy_names = dict(
            (id(v), k) for k, v in list(self._needs.items()))

        for n in self.dependencies:
            n.add_listener(self)

        self._finalized_dependencies = set()
        self._enqueued_dependencies = set()

    @property
    def dependencies(self):
        if isinstance(self._needs, dict):
            return list(self._needs.values())
        return self._needs

    def _dependency_name(self, pusher):
        return self._dependecy_names[id(pusher)]

    def __repr__(self):
        return self.__class__.__name__

    def __str__(self):
        return self.__repr__()

    def __enter__(self):
        return self

    def __exit__(self, t, value, traceback):
        pass

    @property
    def version(self):
        return self.__class__.__name__

    @property
    def needs(self):
        return self._needs

    @property
    def dependency_count(self):
        return len(self._needs)

    @property
    def is_root(self):
        return not self._needs

    @property
    def is_leaf(self):
        return not self._listeners

    def add_listener(self, listener):
        self._listeners.append(listener)

    def find_listener(self, predicate):
        for l in self._listeners:
            return l if predicate(l) else l.find_listener(predicate)
        return None

    def disconnect(self):
        for e in self.dependencies:
            try:
                e._listeners.remove(self)
            except ValueError:
                pass

    def _enqueue(self, data, pusher):
        self._cache = data

    def _dequeue(self):
        if self._cache is None:
            raise NotEnoughData()

        v, self._cache = self._cache, None
        return v

    def _process(self, data):
        yield data

    def _first_chunk(self, data):
        return data

    def _last_chunk(self):
        return iter(())

    def _finalize(self, pusher):
        pass

    @property
    def _finalized(self):
        """
        Return true if all dependencies have informed this node that they'll
        be sending no more data (by calling _finalize()), and that they have
        sent at least one batch of data (by calling enqueue())
        """
        return \
            len(self._finalized_dependencies) >= self.dependency_count \
            and len(self._enqueued_dependencies) >= self.dependency_count

    def _push(self, data, queue=None):
        queue.appendleft((
            id(self),
            self.process.__name__,
            {'pusher': self, 'data': data, 'queue': queue}))

    def _finish(self, pusher=None, queue=None):
        self._finalize(pusher)

        if pusher in self.dependencies:
            self._finalized_dependencies.add(id(pusher))

        if pusher:
            return
        queue.appendleft((
            id(self),
            self._finish.__name__,
            {'pusher': self, 'queue': queue}))

    def process(self, data=None, pusher=None, queue=None):
        if data is not None:
            self._enqueued_dependencies.add(id(pusher))
            self._enqueue(data, pusher)

        try:
            inp = self._dequeue()
            inp = self._first_chunk(inp)
            self._first_chunk = lambda x: x
            for d in self._process(inp):
                yield self._push(d, queue=queue)
        except NotEnoughData:
            yield None

        if self.is_root or self._finalized:
            for chunk in self._last_chunk():
                yield self._push(chunk, queue=queue)
            self._finish(pusher=None, queue=queue)
            self._push(None, queue=queue)
            yield None 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, **kwargs):
        super(KeySelector, self).__init__(needs=needs, **kwargs)
        self.aspect_key = aspect_key

    def _enqueue(self, data, pusher):
        try:
            # the data is dictionary-like, and contains the aspect key
            data = data[self.aspect_key]
        except TypeError:
            # the data is not a dictionary
            pass
        except KeyError:
            # the data is dictionary-like, but doesn't contain the aspect key
            pass

        super(KeySelector, self)._enqueue(data, pusher) def dictify(x, key_selector=lambda item: id(item)):
    if x is None:
        return OrderedDict()

    if isinstance(x, dict):
        return x

    try:
        return OrderedDict((key_selector(z), z) for z in x)
    except TypeError:
        return OrderedDict({key_selector(x): x})
from .var import Var
from collections import OrderedDict
from .datawriter import DataWriter, BytesIODataWriter, ClobberDataWriter
from .decoder import JSONDecoder, Decoder, GreedyTextDecoder, DecoderNode, \
    BZ2Decoder, PickleDecoder
from .encoder import IdentityEncoder, JSONEncoder, TextEncoder, BZ2Encoder, \
    PickleEncoder
from .extractor import Graph, FunctionalNode, Node, KeySelector
from .util import dictify
import inspect
            return False

        return all(
            [n._can_compute(_id, persistence) for n in self.dependencies])

    def _compute(self, _id=None, persistence=None):
        graph, stream = self._build_partial(_id, persistence)

        kwargs = dict()
        for k, extractor in list(graph.roots().items()):
            try:
                kwargs[k] = extractor._reader
            except AttributeError:
                kwargs[k] = self.reader(_id, k, persistence)

        graph.process(**kwargs)
        return stream

    def compute(self, _id, persistence):
        self._compute(_id, persistence)

    def __call__(self, _id=None, decoder=None, persistence=None):
        if decoder is None:
            decoder = self.decoder

        try:
            raw = self.reader(_id, self.key, persistence)
            decoded = decoder(raw)
            return decoded
        except KeyError:
            pass

        if not self._can_compute(_id, persistence):
            raise AttributeError('%s cannot be computed' % self.key)

        stream = self._compute(_id, persistence)

        if stream is None:
            stream = self.reader(_id, self.key, persistence)

        stream.seek(0)
        decoded = decoder(stream)
        return decoded

    def _build_partial(self, _id, persistence):
        features = self._partial(_id, persistence=persistence)
        g = Graph()
        for feat in list(features.values()):
            e = feat._build_extractor(_id, g, persistence)
            if feat.key == self.key:
                stream = e.find_listener(
                    lambda x: isinstance(x, BytesIODataWriter))
                if stream is not None:
                    stream = stream._stream

        return g, stream

    def _partial(self, _id, features=None, persistence=None):
        """
        TODO: _partial is a shit name for this, kind of.  I'm building a graph
        such that I can only do work necessary to compute self, and no more
        """

        root = features is None

        stored = self._stored(_id, persistence)
        is_cached = self.store and stored

        if self.store and not stored:
            data_writer = None
        elif root:
            data_writer = BytesIODataWriter
        else:
            data_writer = None

        should_store = self.store and not stored
        nf = self.copy(
            extractor=DecoderNode if is_cached else self.extractor,
            store=root or should_store,
            needs=dict(),
            data_writer=data_writer,
            persistence=self.persistence,
            extractor_args=dict(decodifier=self.decoder, version=self.version) \
                if is_cached else self.extractor_args)

        if root:
            features = dict()

        features[self.key] = nf

        if not is_cached:
            for k, v in list(self.needs.items()):
                v._partial(_id, features=features, persistence=persistence)
                nf.needs[k] = features[v.key]

        return features

    def _depends_on(self, _id, graph, persistence, **kwargs):
        needs = OrderedDict()

        for k, f in list(self.needs.items()):
            if f.key in graph:
                needs[k] = graph[f.key]
                continue

            e = f._build_extractor(_id, graph, persistence, **kwargs)
            needs[k] = e

        return needs

    def _build_extractor(self, _id, graph, persistence, **kwargs):
        try:
            return graph[self.key]
        except KeyError:
            pass

        needs = self._depends_on(_id, graph, persistence, **kwargs)

        extractor_args = dict()
        for k, v in list(self.extractor_args.items()):

            if isinstance(v, Var) and v.name == k: