Basic styling, lots of feature work

app/timer/models/locations.py

@@ -3,6 +3,16 @@ from .utils import InheritanceQuerySet
 from .owners import Corporation
 
 class LocationManager(models.Manager):
+    use_for_related_fields = True
+
+    def get_query_set(self):
+        return InheritanceQuerySet(self.model)
+
+    def select_subclasses(self, *subclasses):
+        return self.get_query_set().select_subclasses(*subclasses)
+
+    def get_subclass(self, *args, **kwargs):
+        return self.get_query_set().select_subclasses().get(*args, **kwargs)
 
     def all_subclassed(self):
         return InheritanceQuerySet(model=self.model).select_subclasses()
@@ -15,8 +25,17 @@ class Location(models.Model):
     y = models.BigIntegerField('Y Location', null=True)
     z = models.BigIntegerField('Z Location', null=True)
 
+    objects = LocationManager()
+
     def __unicode__(self):
-        return "%(name)s (%(id)d)" % self.__dict__
+        return "%(name)s" % self.__dict__
+
+    def get_subclass(self):
+        return Location.objects.get_subclass(pk=self.pk)
+
+    @property
+    def get_type(self):
+        return self.get_subclass().__class__.__name__
 
     class Meta:
         app_label = 'timer'
@@ -77,6 +96,10 @@ class Planet(Location):
 class Moon(Location):
     planet = models.ForeignKey(Planet, related_name='moons')
 
+    @property
+    def system(self):
+        return self.planet.system
+
     class Meta:
         app_label = 'timer'
 

app/timer/models/timer.py

@@ -12,11 +12,11 @@ class Timer(models.Model):
         (STATE_EXPIRED, 'Expired'),
     )
 
-    TYPE_SHEILD_REENFORCEMENT = 1
+    TYPE_SHIELD_REENFORCEMENT = 1
     TYPE_ARMOR_REENFORCEMENT = 2
 
     TYPE_CHOICES = (
-        (TYPE_SHEILD_REENFORCEMENT, 'Sheild Reenforcement'),
+        (TYPE_SHIELD_REENFORCEMENT, 'Shield Reenforcement'),
         (TYPE_ARMOR_REENFORCEMENT, 'Armor Reenforcement'),
     )
 
@@ -30,6 +30,12 @@ class Timer(models.Model):
             return self.STATE_EXPIRED
         return self.STATE_ACTIVE
 
+    def get_state_display(self):
+        state = self.state
+        for v, disp in self.STATE_CHOICES:
+            if state == v:
+                return disp
+
     class Meta:
         app_label = 'timer'
         ordering = ['-expiration']

app/timer/models/utils.py

@@ -1,30 +1,91 @@
-from django.db.models.fields.related import SingleRelatedObjectDescriptor
+import django
 from django.db.models.query import QuerySet
+from django.db.models.fields.related import OneToOneField
+from django.core.exceptions import ObjectDoesNotExist
+
+try:
+    from django.db.models.constants import LOOKUP_SEP
+except ImportError: # Django < 1.5
+    from django.db.models.sql.constants import LOOKUP_SEP
 
 
 class InheritanceQuerySet(QuerySet):
     def select_subclasses(self, *subclasses):
         if not subclasses:
-            subclasses = [o for o in dir(self.model)
-                          if isinstance(getattr(self.model, o), SingleRelatedObjectDescriptor) \
-                          and issubclass(getattr(self.model, o).related.model, self.model)]
+            # only recurse one level on Django < 1.6 to avoid triggering
+            # https://code.djangoproject.com/ticket/16572
+            levels = None
+            if django.VERSION < (1, 6, 0):
+                levels = 1
+            subclasses = self._get_subclasses_recurse(self.model, levels=levels)
+            # workaround https://code.djangoproject.com/ticket/16855
+        field_dict = self.query.select_related
         new_qs = self.select_related(*subclasses)
+        if isinstance(new_qs.query.select_related, dict) and isinstance(field_dict, dict):
+            new_qs.query.select_related.update(field_dict)
         new_qs.subclasses = subclasses
         return new_qs
 
+
     def _clone(self, klass=None, setup=False, **kwargs):
-        try:
-            kwargs.update({'subclasses': self.subclasses})
-        except AttributeError:
-            pass
+        for name in ['subclasses', '_annotated']:
+            if hasattr(self, name):
+                kwargs[name] = getattr(self, name)
         return super(InheritanceQuerySet, self)._clone(klass, setup, **kwargs)
 
+
+    def annotate(self, *args, **kwargs):
+        qset = super(InheritanceQuerySet, self).annotate(*args, **kwargs)
+        qset._annotated = [a.default_alias for a in args] + list(kwargs.keys())
+        return qset
+
+
     def iterator(self):
         iter = super(InheritanceQuerySet, self).iterator()
         if getattr(self, 'subclasses', False):
             for obj in iter:
-                obj = [getattr(obj, s) for s in self.subclasses if getattr(obj, s)] or [obj]
-                yield obj[0]
+                sub_obj = None
+                for s in self.subclasses:
+                    sub_obj = self._get_sub_obj_recurse(obj, s)
+                    if sub_obj:
+                        break
+                if not sub_obj:
+                    sub_obj = obj
+
+                if getattr(self, '_annotated', False):
+                    for k in self._annotated:
+                        setattr(sub_obj, k, getattr(obj, k))
+
+                yield sub_obj
         else:
             for obj in iter:
-                yield obj
+                yield obj
+
+
+    def _get_subclasses_recurse(self, model, levels=None):
+        rels = [rel for rel in model._meta.get_all_related_objects()
+                if isinstance(rel.field, OneToOneField)
+            and issubclass(rel.field.model, model)]
+        subclasses = []
+        if levels:
+            levels -= 1
+        for rel in rels:
+            if levels or levels is None:
+                for subclass in self._get_subclasses_recurse(
+                        rel.field.model, levels=levels):
+                    subclasses.append(rel.var_name + LOOKUP_SEP + subclass)
+            subclasses.append(rel.var_name)
+        return subclasses
+
+
+    def _get_sub_obj_recurse(self, obj, s):
+        rel, _, s = s.partition(LOOKUP_SEP)
+        try:
+            node = getattr(obj, rel)
+        except ObjectDoesNotExist:
+            return None
+        if s:
+            child = self._get_sub_obj_recurse(node, s)
+            return child or node
+        else:
+            return node