# Source code for sqlalchemy_utils.types.pg_composite

    """
    CompositeType provides means to interact with
    `PostgreSQL composite types`_. Currently this type features:

    * Easy attribute access to composite type fields
    * Supports SQLAlchemy TypeDecorator types
    * Ability to include composite types as part of PostgreSQL arrays
    * Type creation and dropping

    Installation
    ^^^^^^^^^^^^

    CompositeType automatically attaches `before_create` and `after_drop` DDL
    listeners. These listeners create and drop the composite type in the
    database. This means it works out of the box in your test environment where
    you create the tables on each test run.

    When you already have your database set up you should call
    :func:`register_composites` after you've set up all models.

    ::

        register_composites(conn)

    Usage
    ^^^^^

    ::

        from collections import OrderedDict

        import sqlalchemy as sa
        from sqlalchemy_utils import CompositeType, CurrencyType

        class Account(Base):
            __tablename__ = 'account'
            id = sa.Column(sa.Integer, primary_key=True)
            balance = sa.Column(
                CompositeType(
                    'money_type',
                    [
                        sa.Column('currency', CurrencyType),
                        sa.Column('amount', sa.Integer)
                    ]
                )
            )

    Accessing fields
    ^^^^^^^^^^^^^^^^

    CompositeType provides attribute access to underlying fields. In the following
    example we find all accounts with balance amount more than 5000.

    ::

        session.query(Account).filter(Account.balance.amount > 5000)

    Arrays of composites
    ^^^^^^^^^^^^^^^^^^^^

    ::

        from sqlalchemy_utils import CompositeArray

        class Account(Base):
            __tablename__ = 'account'
            id = sa.Column(sa.Integer, primary_key=True)
            balances = sa.Column(
                CompositeArray(
                    CompositeType(
                        'money_type',
                        [
                            sa.Column('currency', CurrencyType),
                            sa.Column('amount', sa.Integer)
                        ]
                    )
                )
            )

    .. _PostgreSQL composite types:
        http://www.postgresql.org/docs/current/static/rowtypes.html

    Related links:

    http://schinckel.net/2014/09/24/using-postgres-composite-types-in-django/
    """
    from collections import namedtuple

    import six
    import sqlalchemy as sa
    from sqlalchemy.dialects.postgresql import ARRAY
    from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
    from sqlalchemy.ext.compiler import compiles
    from sqlalchemy.schema import _CreateDropBase
    from sqlalchemy.sql.expression import FunctionElement
    from sqlalchemy.types import (
        SchemaType,
        to_instance,
        TypeDecorator,
        UserDefinedType
    )

    from .. import ImproperlyConfigured

    psycopg2 = None
    CompositeCaster = None
    adapt = None
    AsIs = None
    register_adapter = None
    try:
        import psycopg2
        from psycopg2.extras import CompositeCaster
        from psycopg2.extensions import adapt, AsIs, register_adapter
    except ImportError:
        pass

    class CompositeElement(FunctionElement):
        """
        Instances of this class wrap a Postgres composite type.
        """
        def __init__(self, base, field, type_):
            self.name = field
            self.type = to_instance(type_)

            super(CompositeElement, self).__init__(base)

    @compiles(CompositeElement)
    def _compile_pgelem(expr, compiler, **kw):
        return '(%s).%s' % (compiler.process(expr.clauses, **kw), expr.name)

    class CompositeArray(ARRAY):
        def _proc_array(self, arr, itemproc, dim, collection):
            if dim is None:
                if isinstance(self.item_type, CompositeType):
                    arr = [itemproc(a) for a in arr]
                    return arr
            return ARRAY._proc_array(self, arr, itemproc, dim, collection)

    # TODO: Make the registration work on connection level instead of global level
    registered_composites = {}

    class CompositeType(UserDefinedType, SchemaType):
        """
        Represents a PostgreSQL composite type.

        :param name:
            Name of the composite type.
        :param columns:
            List of columns that this composite type consists of
        """
        python_type = tuple

        class comparator_factory(UserDefinedType.Comparator):
            def __getattr__(self, key):
                try:
                    type_ = self.type.typemap[key]
                except KeyError:
                    raise KeyError(
                        "Type '%s' doesn't have an attribute: '%s'" % (
                            self.name, key
                        )
                    )

                return CompositeElement(self.expr, key, type_)

        def __init__(self, name, columns):
            if psycopg2 is None:
                raise ImproperlyConfigured(
                    "'psycopg2' package is required in order to use CompositeType."
                )
            SchemaType.__init__(self)
            self.name = name
            self.columns = columns
            if name in registered_composites:
                self.type_cls = registered_composites[name].type_cls
            else:
                self.type_cls = namedtuple(
                    self.name, [c.name for c in columns]
                )
            registered_composites[name] = self

            class Caster(CompositeCaster):
                def make(obj, values):
                    return self.type_cls(*values)

            self.caster = Caster
            attach_composite_listeners()

        def get_col_spec(self):
            return self.name

        def bind_processor(self, dialect):
            def process(value):
                if value is None:
                    return None
                processed_value = []
                for i, column in enumerate(self.columns):
                    if isinstance(column.type, TypeDecorator):
                        processed_value.append(
                            column.type.process_bind_param(
                                value[i], dialect
                            )
                        )
                    else:
                        processed_value.append(value[i])
                return self.type_cls(*processed_value)
            return process

        def result_processor(self, dialect, coltype):
            def process(value):
                if value is None:
                    return None
                cls = value.__class__
                kwargs = {}
                for column in self.columns:
                    if isinstance(column.type, TypeDecorator):
                        kwargs[column.name] = column.type.process_result_value(
                            getattr(value, column.name), dialect
                        )
                    else:
                        kwargs[column.name] = getattr(value, column.name)
                return cls(**kwargs)
            return process

        def create(self, bind=None, checkfirst=None):
            if (
                not checkfirst or
                not bind.dialect.has_type(bind, self.name, schema=self.schema)
            ):
                bind.execute(CreateCompositeType(self))

        def drop(self, bind=None, checkfirst=True):
            if (
                checkfirst and
                bind.dialect.has_type(bind, self.name, schema=self.schema)
            ):
                bind.execute(DropCompositeType(self))

    def register_psycopg2_composite(dbapi_connection, composite):
        psycopg2.extras.register_composite(
            composite.name,
            dbapi_connection,
            globally=True,
            factory=composite.caster
        )

        def adapt_composite(value):
            adapted = [
                adapt(
                    getattr(value, column.name)
                    if not isinstance(column.type, TypeDecorator)
                    else column.type.process_bind_param(
                        getattr(value, column.name),
                        PGDialect_psycopg2()
                    )
                )
                for column in
                composite.columns
            ]
            for value in adapted:
                if hasattr(value, 'prepare'):
                    value.prepare(dbapi_connection)
            values = [
                value.getquoted().decode(dbapi_connection.encoding)
                if six.PY3
                else value.getquoted()
                for value in adapted
            ]
            return AsIs("(%s)::%s" % (', '.join(values), composite.name))

        register_adapter(composite.type_cls, adapt_composite)

    def before_create(target, connection, **kw):
        for name, composite in registered_composites.items():
            composite.create(connection, checkfirst=True)
            register_psycopg2_composite(
                connection.connection.connection,
                composite
            )

    def after_drop(target, connection, **kw):
        for name, composite in registered_composites.items():
            composite.drop(connection, checkfirst=True)

    def register_composites(connection):
        for name, composite in registered_composites.items():
            register_psycopg2_composite(
                connection.connection.connection,
                composite
            )

    def attach_composite_listeners():
        listeners = [
            (sa.MetaData, 'before_create', before_create),
            (sa.MetaData, 'after_drop', after_drop),
        ]
        for listener in listeners:
            if not sa.event.contains(*listener):
                sa.event.listen(*listener)

    def remove_composite_listeners():
        listeners = [
            (sa.MetaData, 'before_create', before_create),
            (sa.MetaData, 'after_drop', after_drop),
        ]
        for listener in listeners:
            if sa.event.contains(*listener):
                sa.event.remove(*listener)

    class CreateCompositeType(_CreateDropBase):
        pass

    @compiles(CreateCompositeType)
    def _visit_create_composite_type(create, compiler, **kw):
        type_ = create.element
        fields = ', '.join(
            '{name} {type}'.format(
                name=column.name,
                type=compiler.dialect.type_compiler.process(
                    to_instance(column.type)
                )
            )
            for column in type_.columns
        )

        return 'CREATE TYPE {name} AS ({fields})'.format(
            name=compiler.preparer.format_type(type_),
            fields=fields
        )

    class DropCompositeType(_CreateDropBase):
        pass

    @compiles(DropCompositeType)
    def _visit_drop_composite_type(drop, compiler, **kw):
        type_ = drop.element

        return 'DROP TYPE {name}'.format(name=compiler.preparer.format_type(type_))