Migrate factory.alchemy to class Meta

1 files changed, 18 insertions, 5 deletions

diff --git a/factory/alchemy.py b/factory/alchemy.py
index cec15c9..b956d7e 100644
--- a/factory/alchemy.py
+++ b/factory/alchemy.py
@@ -24,17 +24,30 @@ from sqlalchemy.sql.functions import max
 from . import base
 
 
+class SQLAlchemyOptions(base.FactoryOptions):
+    def _build_default_options(self):
+        return super(SQLAlchemyOptions, self)._build_default_options() + [
+            base.OptionDefault('sqlalchemy_session', None, inherit=True),
+        ]
+
+
 class SQLAlchemyModelFactory(base.Factory):
     """Factory for SQLAlchemy models. """
 
-    ABSTRACT_FACTORY = True
-    FACTORY_SESSION = None
+    _options_class = SQLAlchemyOptions
+    class Meta:
+        abstract = True
+
+    _OLDSTYLE_ATTRIBUTES = base.Factory._OLDSTYLE_ATTRIBUTES.copy()
+    _OLDSTYLE_ATTRIBUTES.update({
+        'FACTORY_SESSION': 'sqlalchemy_session',
+    })
 
     @classmethod
     def _setup_next_sequence(cls, *args, **kwargs):
         """Compute the next available PK, based on the 'pk' database field."""
-        session = cls.FACTORY_SESSION
-        model = cls.FACTORY_FOR
+        session = cls._meta.sqlalchemy_session
+        model = cls._meta.target
         pk = getattr(model, model.__mapper__.primary_key[0].name)
         max_pk = session.query(max(pk)).one()[0]
         if isinstance(max_pk, int):
@@ -45,7 +58,7 @@ class SQLAlchemyModelFactory(base.Factory):
     @classmethod
     def _create(cls, target_class, *args, **kwargs):
         """Create an instance of the model, and save it to the database."""
-        session = cls.FACTORY_SESSION
+        session = cls._meta.sqlalchemy_session
         obj = target_class(*args, **kwargs)
         session.add(obj)
         return obj