from django.db.models import Aggregate
from django.contrib.gis.db.backend import SpatialBackend
from django.contrib.gis.db.models.sql import GeomField

class GeoAggregate(Aggregate):

    def add_to_query(self, query, alias, col, source, is_summary):
        if hasattr(source, 'geom_type'):
            # Doing additional setup on the Query object for spatial aggregates.
            aggregate = getattr(query.aggregates_module, self.name)
            
            # Adding a conversion class instance and any selection wrapping
            # SQL (e.g., needed by Oracle).
            if aggregate.conversion_class is GeomField:
                query.extra_select_fields[alias] = GeomField()
                if SpatialBackend.select:
                    query.custom_select[alias] = SpatialBackend.select
     
        super(GeoAggregate, self).add_to_query(query, alias, col, source, is_summary)

class Collect(GeoAggregate):
    name = 'Collect'

class Extent(GeoAggregate):
    name = 'Extent'

class MakeLine(GeoAggregate):
    name = 'MakeLine'

class Union(GeoAggregate):
    name = 'Union'