Merge pull request #135 from jwhitlock/override_manager_in_register

Allow overriding records manager class in register

Author: treyhunner

simple_history/__init__.py

@@ -3,20 +3,26 @@
 __version__ = '1.5.1'
 
 
-def register(model, app=None, manager_name='history', **records_config):
+def register(
+        model, app=None, manager_name='history', records_class=None,
+        **records_config):
     """
     Create historical model for `model` and attach history manager to `model`.
 
     Keyword arguments:
     app -- App to install historical model into (defaults to model.__module__)
     manager_name -- class attribute name to use for historical manager
+    records_class -- class to use for history relation (defaults to
+        HistoricalRecords)
 
     This method should be used as an alternative to attaching an
     `HistoricalManager` instance directly to `model`.
     """
     from . import models
     if model._meta.db_table not in models.registered_models:
-        records = models.HistoricalRecords(**records_config)
+        if records_class is None:
+            records_class = models.HistoricalRecords
+        records = records_class(**records_config)
         records.manager_name = manager_name
         records.module = app and ("%s.models" % app) or model.__module__
         records.add_extra_methods(model)

simple_history/tests/models.py

@@ -54,6 +54,25 @@ class Choice(models.Model):
 register(Choice)
 
 
+class Voter(models.Model):
+    user = models.ForeignKey(User)
+    choice = models.ForeignKey(Choice, related_name='voters')
+
+
+class HistoricalRecordsVerbose(HistoricalRecords):
+    def get_extra_fields(self, model, fields):
+        def verbose_str(self):
+            return '%s changed by %s as of %s' % (
+                self.history_object, self.history_user, self.history_date)
+
+        extra_fields = super(
+            HistoricalRecordsVerbose, self).get_extra_fields(model, fields)
+        extra_fields['__str__'] = verbose_str
+        return extra_fields
+
+register(Voter, records_class=HistoricalRecordsVerbose)
+
+
 class Place(models.Model):
     name = models.CharField(max_length=100)
 

simple_history/tests/tests/test_models.py

@@ -20,9 +20,9 @@
 from simple_history.models import HistoricalRecords, convert_auto_field
 from simple_history import register
 from ..models import (
-    AdminProfile, Bookcase, MultiOneToOne, Poll, Choice, Restaurant, Person,
-    FileModel, Document, Book, HistoricalPoll, Library, State, AbstractBase,
-    ConcreteAttr, ConcreteUtil, SelfFK, Temperature, WaterLevel,
+    AdminProfile, Bookcase, MultiOneToOne, Poll, Choice, Voter, Restaurant,
+    Person, FileModel, Document, Book, HistoricalPoll, Library, State,
+    AbstractBase, ConcreteAttr, ConcreteUtil, SelfFK, Temperature, WaterLevel,
     ExternalModel1, ExternalModel3, UnicodeVerboseName, HistoricalChoice,
     HistoricalState, HistoricalCustomFKError
 )
@@ -305,6 +305,16 @@ def test_reregister(self):
         self.assertTrue(hasattr(User, 'histories'))
         self.assertFalse(hasattr(User, 'again'))
 
+    def test_register_custome_records(self):
+        self.assertEqual(len(Voter.history.all()), 0)
+        poll = Poll.objects.create(pub_date=today)
+        choice = Choice.objects.create(poll=poll, votes=0)
+        user = User.objects.create(username='voter')
+        voter = Voter.objects.create(choice=choice, user=user)
+        self.assertEqual(len(voter.history.all()), 1)
+        expected = 'Voter object changed by None as of '
+        self.assertEqual(expected, str(voter.history.all()[0])[:len(expected)])
+
 
 class CreateHistoryModelTests(TestCase):