sqlalchemy/test/ext/test_extendedattr.py

import sqlalchemy as sa
from sqlalchemy import Column
from sqlalchemy import event
from sqlalchemy import Integer
from sqlalchemy import Table
from sqlalchemy import util
from sqlalchemy.ext import instrumentation
from sqlalchemy.orm import attributes
from sqlalchemy.orm import class_mapper
from sqlalchemy.orm import clear_mappers
from sqlalchemy.orm import events
from sqlalchemy.orm.attributes import del_attribute
from sqlalchemy.orm.attributes import get_attribute
from sqlalchemy.orm.attributes import set_attribute
from sqlalchemy.orm.instrumentation import is_instrumented
from sqlalchemy.orm.instrumentation import manager_of_class
from sqlalchemy.orm.instrumentation import register_class
from sqlalchemy.testing import assert_raises
from sqlalchemy.testing import assert_raises_message
from sqlalchemy.testing import eq_
from sqlalchemy.testing import expect_raises_message
from sqlalchemy.testing import fixtures
from sqlalchemy.testing import is_
from sqlalchemy.testing import is_not
from sqlalchemy.testing import ne_
from sqlalchemy.testing.util import decorator


def _register_attribute(class_, key, **kw):
    kw.setdefault("comparator", object())
    kw.setdefault("parententity", object())

    attributes._register_attribute(class_, key, **kw)


@decorator
def modifies_instrumentation_finders(fn, *args, **kw):
    pristine = instrumentation.instrumentation_finders[:]
    try:
        fn(*args, **kw)
    finally:
        del instrumentation.instrumentation_finders[:]
        instrumentation.instrumentation_finders.extend(pristine)


class _ExtBase:
    @classmethod
    def teardown_test_class(cls):
        instrumentation._reinstall_default_lookups()


class MyTypesManager(instrumentation.InstrumentationManager):
    def instrument_attribute(self, class_, key, attr):
        pass

    def install_descriptor(self, class_, key, attr):
        pass

    def uninstall_descriptor(self, class_, key):
        pass

    def instrument_collection_class(self, class_, key, collection_class):
        return MyListLike

    def get_instance_dict(self, class_, instance):
        return instance._goofy_dict

    def initialize_instance_dict(self, class_, instance):
        instance.__dict__["_goofy_dict"] = {}

    def install_state(self, class_, instance, state):
        instance.__dict__["_my_state"] = state

    def state_getter(self, class_):
        return lambda instance: instance.__dict__["_my_state"]


class MyListLike(list):
    # add @appender, @remover decorators as needed
    _sa_iterator = list.__iter__
    _sa_linker = None

    def _sa_appender(self, item, _sa_initiator=None):
        if _sa_initiator is not False:
            self._sa_adapter.fire_append_event(item, _sa_initiator)
        list.append(self, item)

    append = _sa_appender

    def _sa_remover(self, item, _sa_initiator=None):
        self._sa_adapter.fire_pre_remove_event(_sa_initiator)
        if _sa_initiator is not False:
            self._sa_adapter.fire_remove_event(item, _sa_initiator)
        list.remove(self, item)

    remove = _sa_remover


MyBaseClass, MyClass = None, None


class DisposeTest(_ExtBase, fixtures.TestBase):
    def test_unregister(self, registry):
        class MyClassState(instrumentation.InstrumentationManager):
            def manage(self, class_, manager):
                setattr(class_, "xyz", manager)

            def unregister(self, class_, manager):
                delattr(class_, "xyz")

            def manager_getter(self, class_):
                def get(cls):
                    return cls.xyz

                return get

        class MyClass:
            __sa_instrumentation_manager__ = MyClassState

        assert attributes.opt_manager_of_class(MyClass) is None

        with expect_raises_message(
            sa.orm.exc.UnmappedClassError,
            r"Can't locate an instrumentation manager for class .*MyClass",
        ):
            attributes.manager_of_class(MyClass)

        t = Table(
            "my_table",
            registry.metadata,
            Column("id", Integer, primary_key=True),
        )

        registry.map_imperatively(MyClass, t)

        manager = attributes.opt_manager_of_class(MyClass)
        is_not(manager, None)
        is_(manager, MyClass.xyz)

        registry.configure()

        registry.dispose()

        manager = attributes.opt_manager_of_class(MyClass)
        is_(manager, None)

        assert not hasattr(MyClass, "xyz")


class UserDefinedExtensionTest(_ExtBase, fixtures.ORMTest):
    @classmethod
    def setup_test_class(cls):
        global MyBaseClass, MyClass

        class MyBaseClass:
            __sa_instrumentation_manager__ = (
                instrumentation.InstrumentationManager
            )

        class MyClass:
            # This proves that a staticmethod will work here; don't
            # flatten this back to a class assignment!
            def __sa_instrumentation_manager__(cls):
                return MyTypesManager(cls)

            __sa_instrumentation_manager__ = staticmethod(
                __sa_instrumentation_manager__
            )

            # This proves SA can handle a class with non-string dict keys
            # Since python 3.13 non-string key raise a runtime warning.
            if util.cpython and not util.py313:
                locals()[42] = 99  # Don't remove this line!

            def __init__(self, **kwargs):
                for k in kwargs:
                    setattr(self, k, kwargs[k])

            def __getattr__(self, key):
                if is_instrumented(self, key):
                    return get_attribute(self, key)
                else:
                    try:
                        return self._goofy_dict[key]
                    except KeyError:
                        raise AttributeError(key)

            def __setattr__(self, key, value):
                if is_instrumented(self, key):
                    set_attribute(self, key, value)
                else:
                    self._goofy_dict[key] = value

            def __hasattr__(self, key):
                if is_instrumented(self, key):
                    return True
                else:
                    return key in self._goofy_dict

            def __delattr__(self, key):
                if is_instrumented(self, key):
                    del_attribute(self, key)
                else:
                    del self._goofy_dict[key]

    def teardown_test(self):
        clear_mappers()

    def test_instance_dict(self):
        class User(MyClass):
            pass

        register_class(User)
        _register_attribute(User, "user_id", uselist=False, useobject=False)
        _register_attribute(User, "user_name", uselist=False, useobject=False)
        _register_attribute(
            User, "email_address", uselist=False, useobject=False
        )

        u = User()
        u.user_id = 7
        u.user_name = "john"
        u.email_address = "lala@123.com"
        eq_(
            u.__dict__,
            {
                "_my_state": u._my_state,
                "_goofy_dict": {
                    "user_id": 7,
                    "user_name": "john",
                    "email_address": "lala@123.com",
                },
            },
        )

    def test_basic(self):
        for base in (object, MyBaseClass, MyClass):

            class User(base):
                pass

            register_class(User)
            _register_attribute(
                User, "user_id", uselist=False, useobject=False
            )
            _register_attribute(
                User, "user_name", uselist=False, useobject=False
            )
            _register_attribute(
                User, "email_address", uselist=False, useobject=False
            )

            u = User()
            u.user_id = 7
            u.user_name = "john"
            u.email_address = "lala@123.com"

            eq_(u.user_id, 7)
            eq_(u.user_name, "john")
            eq_(u.email_address, "lala@123.com")
            attributes.instance_state(u)._commit_all(
                attributes.instance_dict(u)
            )
            eq_(u.user_id, 7)
            eq_(u.user_name, "john")
            eq_(u.email_address, "lala@123.com")

            u.user_name = "heythere"
            u.email_address = "foo@bar.com"
            eq_(u.user_id, 7)
            eq_(u.user_name, "heythere")
            eq_(u.email_address, "foo@bar.com")

    def test_deferred(self):
        for base in (object, MyBaseClass, MyClass):

            class Foo(base):
                pass

            data = {"a": "this is a", "b": 12}

            def loader(state, keys, passive):
                for k in keys:
                    state.dict[k] = data[k]
                return attributes.ATTR_WAS_SET

            manager = register_class(Foo)
            manager.expired_attribute_loader = loader
            _register_attribute(Foo, "a", uselist=False, useobject=False)
            _register_attribute(Foo, "b", uselist=False, useobject=False)

            if base is object:
                assert Foo not in (
                    instrumentation._instrumentation_factory._state_finders
                )
            else:
                assert Foo in (
                    instrumentation._instrumentation_factory._state_finders
                )

            f = Foo()
            attributes.instance_state(f)._expire(
                attributes.instance_dict(f), set()
            )
            eq_(f.a, "this is a")
            eq_(f.b, 12)

            f.a = "this is some new a"
            attributes.instance_state(f)._expire(
                attributes.instance_dict(f), set()
            )
            eq_(f.a, "this is a")
            eq_(f.b, 12)

            attributes.instance_state(f)._expire(
                attributes.instance_dict(f), set()
            )
            f.a = "this is another new a"
            eq_(f.a, "this is another new a")
            eq_(f.b, 12)

            attributes.instance_state(f)._expire(
                attributes.instance_dict(f), set()
            )
            eq_(f.a, "this is a")
            eq_(f.b, 12)

            del f.a
            eq_(f.a, None)
            eq_(f.b, 12)

            attributes.instance_state(f)._commit_all(
                attributes.instance_dict(f)
            )
            eq_(f.a, None)
            eq_(f.b, 12)

    def test_inheritance(self):
        """tests that attributes are polymorphic"""

        for base in (object, MyBaseClass, MyClass):

            class Foo(base):
                pass

            class Bar(Foo):
                pass

            register_class(Foo)
            register_class(Bar)

            def func1(state, passive):
                return "this is the foo attr"

            def func2(state, passive):
                return "this is the bar attr"

            def func3(state, passive):
                return "this is the shared attr"

            _register_attribute(
                Foo, "element", uselist=False, callable_=func1, useobject=True
            )
            _register_attribute(
                Foo, "element2", uselist=False, callable_=func3, useobject=True
            )
            _register_attribute(
                Bar, "element", uselist=False, callable_=func2, useobject=True
            )

            x = Foo()
            y = Bar()
            assert x.element == "this is the foo attr"
            assert y.element == "this is the bar attr", y.element
            assert x.element2 == "this is the shared attr"
            assert y.element2 == "this is the shared attr"

    def test_collection_with_backref(self):
        for base in (object, MyBaseClass, MyClass):

            class Post(base):
                pass

            class Blog(base):
                pass

            register_class(Post)
            register_class(Blog)
            _register_attribute(
                Post,
                "blog",
                uselist=False,
                backref="posts",
                trackparent=True,
                useobject=True,
            )
            _register_attribute(
                Blog,
                "posts",
                uselist=True,
                backref="blog",
                trackparent=True,
                useobject=True,
            )
            b = Blog()
            (p1, p2, p3) = (Post(), Post(), Post())
            b.posts.append(p1)
            b.posts.append(p2)
            b.posts.append(p3)
            self.assert_(b.posts == [p1, p2, p3])
            self.assert_(p2.blog is b)

            p3.blog = None
            self.assert_(b.posts == [p1, p2])
            p4 = Post()
            p4.blog = b
            self.assert_(b.posts == [p1, p2, p4])

            p4.blog = b
            p4.blog = b
            self.assert_(b.posts == [p1, p2, p4])

            # assert no failure removing None
            p5 = Post()
            p5.blog = None
            del p5.blog

    def test_history(self):
        for base in (object, MyBaseClass, MyClass):

            class Foo(base):
                pass

            class Bar(base):
                pass

            register_class(Foo)
            register_class(Bar)
            _register_attribute(Foo, "name", uselist=False, useobject=False)
            _register_attribute(
                Foo, "bars", uselist=True, trackparent=True, useobject=True
            )
            _register_attribute(Bar, "name", uselist=False, useobject=False)

            f1 = Foo()
            f1.name = "f1"

            eq_(
                attributes.get_state_history(
                    attributes.instance_state(f1), "name"
                ),
                (["f1"], (), ()),
            )

            b1 = Bar()
            b1.name = "b1"
            f1.bars.append(b1)
            eq_(
                attributes.get_state_history(
                    attributes.instance_state(f1), "bars"
                ),
                ([b1], [], []),
            )

            attributes.instance_state(f1)._commit_all(
                attributes.instance_dict(f1)
            )
            attributes.instance_state(b1)._commit_all(
                attributes.instance_dict(b1)
            )

            eq_(
                attributes.get_state_history(
                    attributes.instance_state(f1), "name"
                ),
                ((), ["f1"], ()),
            )
            eq_(
                attributes.get_state_history(
                    attributes.instance_state(f1), "bars"
                ),
                ((), [b1], ()),
            )

            f1.name = "f1mod"
            b2 = Bar()
            b2.name = "b2"
            f1.bars.append(b2)
            eq_(
                attributes.get_state_history(
                    attributes.instance_state(f1), "name"
                ),
                (["f1mod"], (), ["f1"]),
            )
            eq_(
                attributes.get_state_history(
                    attributes.instance_state(f1), "bars"
                ),
                ([b2], [b1], []),
            )
            f1.bars.remove(b1)
            eq_(
                attributes.get_state_history(
                    attributes.instance_state(f1), "bars"
                ),
                ([b2], [], [b1]),
            )

    def test_null_instrumentation(self):
        class Foo(MyBaseClass):
            pass

        register_class(Foo)
        _register_attribute(Foo, "name", uselist=False, useobject=False)
        _register_attribute(
            Foo, "bars", uselist=True, trackparent=True, useobject=True
        )

        assert Foo.name == attributes.manager_of_class(Foo)["name"]
        assert Foo.bars == attributes.manager_of_class(Foo)["bars"]

    def test_alternate_finders(self):
        """Ensure the generic finder front-end deals with edge cases."""

        class Unknown:
            pass

        class Known(MyBaseClass):
            pass

        register_class(Known)
        k, u = Known(), Unknown()

        assert instrumentation.opt_manager_of_class(Unknown) is None
        assert instrumentation.opt_manager_of_class(Known) is not None
        assert instrumentation.opt_manager_of_class(None) is None

        assert attributes.instance_state(k) is not None
        assert_raises((AttributeError, KeyError), attributes.instance_state, u)
        assert_raises(
            (AttributeError, KeyError), attributes.instance_state, None
        )

    def test_unmapped_not_type_error(self):
        """extension version of the same test in test_mapper.

        fixes #3408
        """
        assert_raises_message(
            sa.exc.ArgumentError,
            "Class object expected, got '5'.",
            class_mapper,
            5,
        )

    def test_unmapped_not_type_error_iter_ok(self):
        """extension version of the same test in test_mapper.

        fixes #3408
        """
        assert_raises_message(
            sa.exc.ArgumentError,
            r"Class object expected, got '\(5, 6\)'.",
            class_mapper,
            (5, 6),
        )


class FinderTest(_ExtBase, fixtures.ORMTest):
    def test_standard(self):
        class A:
            pass

        register_class(A)

        eq_(type(manager_of_class(A)), instrumentation.ClassManager)

    def test_nativeext_interfaceexact(self):
        class A:
            __sa_instrumentation_manager__ = (
                instrumentation.InstrumentationManager
            )

        register_class(A)
        ne_(
            type(attributes.opt_manager_of_class(A)),
            instrumentation.ClassManager,
        )

    def test_nativeext_submanager(self):
        class Mine(instrumentation.ClassManager):
            pass

        class A:
            __sa_instrumentation_manager__ = Mine

        register_class(A)
        eq_(type(manager_of_class(A)), Mine)

    @modifies_instrumentation_finders
    def test_customfinder_greedy(self):
        class Mine(instrumentation.ClassManager):
            pass

        class A:
            pass

        def find(cls):
            return Mine

        instrumentation.instrumentation_finders.insert(0, find)
        register_class(A)
        eq_(type(manager_of_class(A)), Mine)

    @modifies_instrumentation_finders
    def test_customfinder_pass(self):
        class A:
            pass

        def find(cls):
            return None

        instrumentation.instrumentation_finders.insert(0, find)
        register_class(A)

        eq_(type(manager_of_class(A)), instrumentation.ClassManager)


class InstrumentationCollisionTest(_ExtBase, fixtures.ORMTest):
    def test_none(self):
        class A:
            pass

        register_class(A)

        def mgr_factory(cls):
            return instrumentation.ClassManager(cls)

        class B:
            __sa_instrumentation_manager__ = staticmethod(mgr_factory)

        register_class(B)

        class C:
            __sa_instrumentation_manager__ = instrumentation.ClassManager

        register_class(C)

    def test_single_down(self):
        class A:
            pass

        register_class(A)

        def mgr_factory(cls):
            return instrumentation.ClassManager(cls)

        class B(A):
            __sa_instrumentation_manager__ = staticmethod(mgr_factory)

        assert_raises_message(
            TypeError,
            "multiple instrumentation implementations",
            register_class,
            B,
        )

    def test_single_up(self):
        class A:
            pass

        # delay registration

        def mgr_factory(cls):
            return instrumentation.ClassManager(cls)

        class B(A):
            __sa_instrumentation_manager__ = staticmethod(mgr_factory)

        register_class(B)

        assert_raises_message(
            TypeError,
            "multiple instrumentation implementations",
            register_class,
            A,
        )

    def test_diamond_b1(self):
        def mgr_factory(cls):
            return instrumentation.ClassManager(cls)

        class A:
            pass

        class B1(A):
            pass

        class B2(A):
            __sa_instrumentation_manager__ = staticmethod(mgr_factory)

        class C:
            pass

        assert_raises_message(
            TypeError,
            "multiple instrumentation implementations",
            register_class,
            B1,
        )

    def test_diamond_b2(self):
        def mgr_factory(cls):
            return instrumentation.ClassManager(cls)

        class A:
            pass

        class B1(A):
            pass

        class B2(A):
            __sa_instrumentation_manager__ = staticmethod(mgr_factory)

        class C:
            pass

        register_class(B2)
        assert_raises_message(
            TypeError,
            "multiple instrumentation implementations",
            register_class,
            B1,
        )

    def test_diamond_c_b(self):
        def mgr_factory(cls):
            return instrumentation.ClassManager(cls)

        class A:
            pass

        class B1(A):
            pass

        class B2(A):
            __sa_instrumentation_manager__ = staticmethod(mgr_factory)

        class C:
            pass

        register_class(C)

        assert_raises_message(
            TypeError,
            "multiple instrumentation implementations",
            register_class,
            B1,
        )


class ExtendedEventsTest(_ExtBase, fixtures.ORMTest):
    """Allow custom Events implementations."""

    @modifies_instrumentation_finders
    def test_subclassed(self):
        class MyEvents(events.InstanceEvents):
            pass

        class MyClassManager(instrumentation.ClassManager):
            dispatch = event.dispatcher(MyEvents)

        instrumentation.instrumentation_finders.insert(
            0, lambda cls: MyClassManager
        )

        class A:
            pass

        register_class(A)
        manager = instrumentation.manager_of_class(A)
        assert issubclass(manager.dispatch._events, MyEvents)