From 88eeb3689d59c072686df3c17ad0e0884980249a Mon Sep 17 00:00:00 2001 From: Scott Sanderson Date: Wed, 6 Apr 2016 13:25:56 -0400 Subject: [PATCH] MAINT: Make dependencies a dict. --- zipline/pipeline/data/dataset.py | 1 - zipline/pipeline/graph.py | 35 ++++++++------------------------ zipline/pipeline/term.py | 25 +++++++++++++++-------- 3 files changed, 24 insertions(+), 37 deletions(-) diff --git a/zipline/pipeline/data/dataset.py b/zipline/pipeline/data/dataset.py index da236541..b5980d83 100644 --- a/zipline/pipeline/data/dataset.py +++ b/zipline/pipeline/data/dataset.py @@ -114,7 +114,6 @@ class BoundColumn(LoadableTerm): The name of this column. """ mask = AssetExists() - extra_input_rows = 0 inputs = () def __new__(cls, dtype, missing_value, dataset, name): diff --git a/zipline/pipeline/graph.py b/zipline/pipeline/graph.py index 5f815065..6290c937 100644 --- a/zipline/pipeline/graph.py +++ b/zipline/pipeline/graph.py @@ -104,9 +104,9 @@ class TermGraph(DiGraph): zipline.pipeline.engine.SimplePipelineEngine._inputs_for_term zipline.pipeline.engine.SimplePipelineEngine._mask_and_dates_for_term """ - return {(term, dep): self.extra_rows[dep] - term.extra_input_rows + return {(term, dep): self.extra_rows[dep] - additional_extra_rows for term in self - for dep in term.dependencies} + for dep, additional_extra_rows in term.dependencies.items()} @lazyval def extra_rows(self): @@ -119,10 +119,9 @@ class TermGraph(DiGraph): Notes ---- This value depends on the other terms in the graph that require `term` - **as an input**. This is not to be confused with - `term.extra_input_rows`, which is how many extra rows of `term`'s - inputs we need to load, and which is determined entirely by `Term` - itself. + **as an input**. This is not to be confused with `term.dependencies`, + which describes how many additional rows of `term`'s inputs we need to + load, and which is determined entirely by `Term` itself. Example ------- @@ -144,7 +143,7 @@ class TermGraph(DiGraph): See Also -------- zipline.pipeline.graph.TermGraph.offset - zipline.pipeline.term.Term.extra_input_rows + zipline.pipeline.term.Term.dependencies """ return { term: attrs['extra_rows'] @@ -187,30 +186,12 @@ class TermGraph(DiGraph): # Make sure we're going to compute at least `extra_rows` of `term`. self._ensure_extra_rows(term, extra_rows) - # Number of extra rows we need to compute for this term's dependencies. - dependency_extra_rows = extra_rows + term.extra_input_rows - - if isinstance(term, ComputableTerm): - # For computable terms, we want to manually add the term's mask to - # the graph with zero extra rows. A computable term does not - # directly require its mask to have any extra rows. Only loadable - # terms should dictate how many extra rows a mask should compute. - self._add_to_graph( - term.mask, - parents, - extra_rows=0, - ) - self.add_edge(term.mask, term) - dependencies = term.inputs - else: - dependencies = term.dependencies - # Recursively add dependencies. - for dependency in dependencies: + for dependency, additional_extra_rows in term.dependencies.items(): self._add_to_graph( dependency, parents, - extra_rows=dependency_extra_rows, + extra_rows=extra_rows + additional_extra_rows, ) self.add_edge(dependency, term) diff --git a/zipline/pipeline/term.py b/zipline/pipeline/term.py index 07b0ef8f..16faf227 100644 --- a/zipline/pipeline/term.py +++ b/zipline/pipeline/term.py @@ -294,13 +294,13 @@ class Term(with_metaclass(ABCMeta, object)): """ raise NotImplementedError('mask') - @lazyval + @abstractproperty def dependencies(self): """ - A tuple containing all terms that must be computed before this term can - be loaded or computed. + A dictionary mapping terms that must be computed before `self` to the + number of extra rows needed for those terms. """ - return self.inputs + (self.mask,) + raise NotImplementedError('dependencies') class AssetExists(Term): @@ -319,9 +319,8 @@ class AssetExists(Term): """ dtype = bool_dtype dataset = None - extra_input_rows = 0 inputs = () - dependencies = () + dependencies = {} mask = None windowed = False @@ -335,9 +334,12 @@ class LoadableTerm(Term): This is the base class for :class:`zipline.pipeline.data.BoundColumn`. """ - inputs = () windowed = False + @lazyval + def dependencies(self): + return {self.mask: 0} + class ComputableTerm(Term): """ @@ -442,12 +444,17 @@ class ComputableTerm(Term): ) @lazyval - def extra_input_rows(self): + def dependencies(self): """ The number of extra rows needed for each of our inputs to compute this term. """ - return max(0, self.window_length - 1) + extra_input_rows = max(0, self.window_length - 1) + out = {} + for term in self.inputs: + out[term] = extra_input_rows + out[self.mask] = 0 + return out def __repr__(self): return (