import pghistory from django.contrib.postgres.fields import ArrayField from django.db import models from django.urls import reverse from django.utils.text import slugify from core.history import HistoricalSlug from core.models import TrackedModel @pghistory.track() class Company(TrackedModel): class CompanyRole(models.TextChoices): MANUFACTURER = 'MANUFACTURER', 'Ride Manufacturer' DESIGNER = 'DESIGNER', 'Ride Designer' OPERATOR = 'OPERATOR', 'Park Operator' PROPERTY_OWNER = 'PROPERTY_OWNER', 'Property Owner' name = models.CharField(max_length=255) slug = models.SlugField(max_length=255, unique=True) roles = ArrayField( models.CharField(max_length=20, choices=CompanyRole.choices), default=list, blank=True ) description = models.TextField(blank=True) website = models.URLField(blank=True) # General company info founded_date = models.DateField(null=True, blank=True) # Manufacturer-specific fields rides_count = models.IntegerField(default=0) coasters_count = models.IntegerField(default=0) def __str__(self): return self.name def save(self, *args, **kwargs): if not self.slug: self.slug = slugify(self.name) super().save(*args, **kwargs) def get_absolute_url(self): # This will need to be updated to handle different roles return reverse('companies:detail', kwargs={'slug': self.slug}) return '#' @classmethod def get_by_slug(cls, slug): """Get company by current or historical slug""" try: return cls.objects.get(slug=slug), False except cls.DoesNotExist: # Check pghistory first history_model = cls.get_history_model() history_entry = ( history_model.objects.filter(slug=slug) .order_by('-pgh_created_at') .first() ) if history_entry: return cls.objects.get(id=history_entry.pgh_obj_id), True # Check manual slug history as fallback try: historical = HistoricalSlug.objects.get( content_type__model='company', slug=slug ) return cls.objects.get(pk=historical.object_id), True except (HistoricalSlug.DoesNotExist, cls.DoesNotExist): raise cls.DoesNotExist("No company found with this slug") class Meta: ordering = ['name'] verbose_name_plural = 'Companies'