MAINT: Refactor commission model class hierarchies

tests/finance/test_commissions.py

@@ -5,8 +5,12 @@ from nose_parameterized import parameterized
 from pandas import DataFrame
 
 from zipline import TradingAlgorithm
+from zipline.assets import Equity, Future
 from zipline.errors import IncompatibleCommissionModel
 from zipline.finance.commission import (
+    CommissionModel,
+    EquityCommissionModel,
+    FutureCommissionModel,
     PerContract,
     PerDollar,
     PerFutureTrade,
@@ -79,6 +83,47 @@ class CommissionUnitTests(WithAssetFinder, ZiplineTestCase):
         self.assertEqual(0, model.calculate(order, txns[1]))
         self.assertEqual(0, model.calculate(order, txns[2]))
 
+    def test_allowed_asset_types(self):
+        # Custom equities model.
+        class MyEquitiesModel(EquityCommissionModel):
+            def calculate(self, order, transaction):
+                return 0
+
+        self.assertEqual(MyEquitiesModel.allowed_asset_types, (Equity,))
+
+        # Custom futures model.
+        class MyFuturesModel(FutureCommissionModel):
+            def calculate(self, order, transaction):
+                return 0
+
+        self.assertEqual(MyFuturesModel.allowed_asset_types, (Future,))
+
+        # Custom model for both equities and futures.
+        class MyMixedModel(EquityCommissionModel, FutureCommissionModel):
+            def calculate(self, order, transaction):
+                return 0
+
+        self.assertEqual(MyMixedModel.allowed_asset_types, (Equity, Future))
+
+        # Equivalent custom model for both equities and futures.
+        class MyMixedModel(CommissionModel):
+            def calculate(self, order, transaction):
+                return 0
+
+        self.assertEqual(MyMixedModel.allowed_asset_types, (Equity, Future))
+
+        SomeType = type('SomeType', (object,), {})
+
+        # A custom model that defines its own allowed types should take
+        # precedence over the parent class definitions.
+        class MyCustomModel(EquityCommissionModel, FutureCommissionModel):
+            allowed_asset_types = (SomeType,)
+
+            def calculate(self, order, transaction):
+                return 0
+
+        self.assertEqual(MyCustomModel.allowed_asset_types, (SomeType,))
+
     def test_per_trade(self):
         # Test per trade model for equities.
         model = PerTrade(cost=10)

zipline/finance/commission.py

@@ -12,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import abc
 from abc import abstractmethod
 from collections import defaultdict
 
@@ -21,6 +20,7 @@ from toolz import merge
 
 from zipline.assets import Equity, Future
 from zipline.finance.constants import FUTURE_EXCHANGE_FEES_BY_SYMBOL
+from zipline.finance.shared import AllowedAssetMarker, FinancialModelMeta
 from zipline.utils.dummy import DummyMapping
 
 DEFAULT_PER_SHARE_COST = 0.0075              # 0.75 cents per share
@@ -30,7 +30,7 @@ DEFAULT_MINIMUM_COST_PER_EQUITY_TRADE = 1.0  # $1 per trade
 DEFAULT_MINIMUM_COST_PER_FUTURE_TRADE = 1.0  # $1 per trade
 
 
-class CommissionModel(with_metaclass(abc.ABCMeta)):
+class CommissionModel(with_metaclass(FinancialModelMeta)):
     """
     Abstract commission model interface.
 
@@ -70,14 +70,16 @@ class CommissionModel(with_metaclass(abc.ABCMeta)):
         raise NotImplementedError('calculate')
 
 
-class EquityCommissionModel(CommissionModel):
+class EquityCommissionModel(with_metaclass(AllowedAssetMarker,
+                                           CommissionModel)):
     """
     Base class for commission models which only support equities.
     """
     allowed_asset_types = (Equity,)
 
 
-class FutureCommissionModel(CommissionModel):
+class FutureCommissionModel(with_metaclass(AllowedAssetMarker,
+                                           CommissionModel)):
     """
     Base class for commission models which only support futures.
     """
@@ -253,7 +255,7 @@ class PerContract(FutureCommissionModel):
         )
 
 
-class PerEquityTrade(EquityCommissionModel):
+class PerTrade(CommissionModel):
     """
     Calculates a commission for a transaction based on a per trade cost.
 
@@ -352,7 +354,3 @@ class PerDollar(EquityCommissionModel):
         """
         cost_per_share = transaction.price * self.cost_per_dollar
         return abs(transaction.amount) * cost_per_share
-
-
-# Alias PerTrade for backwards compatibility.
-PerTrade = PerEquityTrade

zipline/finance/shared.py

@@ -0,0 +1,44 @@
+
+from abc import ABCMeta
+from itertools import chain
+
+
+class FinancialModelMeta(ABCMeta):
+    """
+    This metaclass allows users to create custom slippage and commission models
+    that support both equities and futures by subclassing the appropriate
+    individualized classes.
+
+    Take for example the following custom model, which subclasses both
+    EquitySlippageModel and FutureSlippageModel together:
+
+        class MyCustomSlippage(EquitySlippageModel, FutureSlippageModel):
+            def process_order(self, data, order):
+                ...
+
+    Ordinarily the first parent class in the MRO ('EquitySlippageModel' in this
+    example) would override the 'allowed_asset_types' class attribute to only
+    include equities. However, this is probably not the expected behavior for a
+    reasonable user, so the metaclass will handle this specific case by
+    manually allowing both equities and futures for the class being created.
+    """
+
+    def __new__(mcls, name, bases, dict_):
+        if 'allowed_asset_types' not in dict_:
+            allowed_asset_types = tuple(
+                chain.from_iterable(
+                    marker.allowed_asset_types
+                    for marker in bases
+                    if isinstance(marker, AllowedAssetMarker)
+                )
+            )
+            if allowed_asset_types:
+                dict_['allowed_asset_types'] = allowed_asset_types
+
+        return super(FinancialModelMeta, mcls).__new__(
+            mcls, name, bases, dict_,
+        )
+
+
+class AllowedAssetMarker(FinancialModelMeta):
+    pass

zipline/finance/slippage.py

@@ -14,8 +14,7 @@
 # limitations under the License.
 from __future__ import division
 
-from abc import ABCMeta, abstractmethod
-from itertools import chain
+from abc import abstractmethod
 import math
 
 import numpy as np
@@ -26,6 +25,7 @@ from toolz import merge
 from zipline.assets import Equity, Future
 from zipline.errors import HistoryWindowStartsBeforeData
 from zipline.finance.constants import ROOT_SYMBOL_TO_ETA
+from zipline.finance.shared import AllowedAssetMarker, FinancialModelMeta
 from zipline.finance.transaction import create_transaction
 from zipline.utils.cache import ExpiringCache
 from zipline.utils.dummy import DummyMapping
@@ -78,45 +78,7 @@ def fill_price_worse_than_limit_price(fill_price, order):
     return False
 
 
-class SlippageModelMeta(ABCMeta):
-    """
-    This metaclass allows users to create custom slippage models that support
-    both equities and futures by subclassing EquityFutureModel and
-    FutureSlippageModel together.
-
-    Take for example the following custom model:
-
-        class MyCustomSlippage(EquitySlippageModel, FutureSlippageModel):
-            def process_order(self, data, order):
-                ...
-
-    Ordinarily the first parent class in the MRO ('EquitySlippageModel' in this
-    example) would override the 'allowed_asset_types' class attribute to only
-    include equities. However, this is probably not the expected behavior for a
-    reasonable user, so the metaclass will handle this specific case by
-    manually allowing both equities and futures for the class being created.
-    """
-
-    def __new__(mcls, name, bases, dict_):
-        if 'allowed_asset_types' not in dict_:
-            allowed_asset_types = tuple(
-                chain.from_iterable(
-                    marker.allowed_asset_types
-                    for marker in bases
-                    if isinstance(marker, AssetSlippageMarker)
-                )
-            )
-            if allowed_asset_types:
-                dict_['allowed_asset_types'] = allowed_asset_types
-
-        return super(SlippageModelMeta, mcls).__new__(mcls, name, bases, dict_)
-
-
-class AssetSlippageMarker(SlippageModelMeta):
-    pass
-
-
-class SlippageModel(with_metaclass(SlippageModelMeta)):
+class SlippageModel(with_metaclass(FinancialModelMeta)):
     """Abstract interface for defining a slippage model.
     """
 
@@ -214,14 +176,14 @@ class SlippageModel(with_metaclass(SlippageModelMeta)):
         return self.__dict__
 
 
-class EquitySlippageModel(with_metaclass(AssetSlippageMarker, SlippageModel)):
+class EquitySlippageModel(with_metaclass(AllowedAssetMarker, SlippageModel)):
     """
     Base class for slippage models which only support equities.
     """
     allowed_asset_types = (Equity,)
 
 
-class FutureSlippageModel(with_metaclass(AssetSlippageMarker, SlippageModel)):
+class FutureSlippageModel(with_metaclass(AllowedAssetMarker, SlippageModel)):
     """
     Base class for slippage models which only support futures.
     """