from copy import copy
try:
    from collections import OrderedDict
except ImportError:
    from ordereddict import OrderedDict

import sqlalchemy as sa
from sqlalchemy_utils import identity


class Operation(object):
    """
    Simple class holding constants for SQL operations.
    """
    INSERT = 0
    UPDATE = 1
    DELETE = 2

    def __init__(self, target, type):
        self.target = target
        self.type = type
        self.processed = False

    def __eq__(self, other):
        return (
            self.target == other.target and
            self.type == other.type
        )

    def __ne__(self, other):
        return not (self == other)


class Operations(object):
    """
    A collection of operations
    """
    def __init__(self):
        self.objects = OrderedDict()

    def format_key(self, target):
        # We cannot use target._sa_instance_state.identity here since object's
        # identity is not yet updated at this phase
        return (target.__class__, identity(target))

    def __contains__(self, target):
        return self.format_key(target) in self.objects

    def __setitem__(self, key, operation):
        self.objects[key] = operation

    def __getitem__(self, key):
        return self.objects[key]

    def __delitem__(self, key):
        del self.objects[key]

    @property
    def entities(self):
        """
        Return a set of changed versioned entities for given session.

        :param session: SQLAlchemy session object
        """
        return set(key[0] for key, _ in self.iteritems())

    def iteritems(self):
        return self.objects.iteritems()

    def items(self):
        return self.objects.items()

    def add(self, operation):
        self[self.format_key(operation.target)] = operation

    def add_insert(self, target):
        if target in self:
            # If the object is deleted and then inserted within the same
            # transaction we are actually dealing with an update.
            self.add(Operation(target, Operation.UPDATE))
        else:
            self.add(Operation(target, Operation.INSERT))

    def add_update(self, target):
        state_copy = copy(sa.inspect(target).committed_state)
        relationships = sa.inspect(target.__class__).relationships
        # Remove all ONETOMANY and MANYTOMANY relationships
        for rel_key, relationship in relationships.items():
            if relationship.direction.name in ['ONETOMANY', 'MANYTOMANY']:
                if rel_key in state_copy:
                    del state_copy[rel_key]

        if state_copy:
            self.add(Operation(target, Operation.UPDATE))

    def add_delete(self, target):
        self.add(Operation(target, Operation.DELETE))


class IdentitySet(object):
    """
    A set of entities

    ::

        identity_set = IdentitySet()
        identity_set.add(User(id=1))

        User(id=1) in identity_set
    """
    def __init__(self):
        self.objects = OrderedDict()

    def format_key(self, obj):
        return (obj.__class__, identity(obj))

    def __contains__(self, obj):
        return self.format_key(obj) in self.objects

    def remove(self, obj):
        del self.objects[self.format_key(obj)]

    def add(self, obj):
        self.objects[self.format_key(obj)] = obj

    def __iter__(self):
        for obj in self.objects:
            yield self.obj

@@ -121,16 +121,27 @@
 edDict%0A%0A
+import six%0A
 import s
@@ -1596,32 +1596,46 @@
         return s
+ix.iteritems(s
 elf.objects.iter
@@ -1629,27 +1629,16 @@
 .objects
-.iteritems(
 )%0A%0A    d

Fix python3 compatibility