Search code examples
pythonclasspython-dataclasses

How to check if a class has an attribute with `'attr' in instance` syntax


I want to be able to use the 'attr' in instance syntax to check if my dataclass has the specified attribute but I can't seem to make it work.

What I want is same behavior as this example with pandas

import pandas as pd

df = pd.DataFrame(columns=['a', 'b', 'c'])
print('a' in df)
True

But just for a custom dataclass

from dataclasses import dataclass

@dataclass
class User:
    email: int
    password: str
    blocked_at: float = None
    
    def __getitem__(self, item):
        return getattr(self, item)

user = User('email@test.com', 'password')

print(user['email'])

'email' in user
email@test.com

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [35], in <cell line: 1>()
----> 1 'email' in user in User.__getitem__(self, item)
      7 def __getitem__(self, item):
----> 8     return getattr(self, item)

TypeError: getattr(): attribute name must be string

Solution

  • What is happening is that you didn't define the correct hook. You want to implement the __contains__ method.

    Because you didn't, the in operator switches to the fallback mode: iterating over the object as if it was a sequence, so it tries object[0],then object[1], etc. until it hits an IndexError or finds something that is equal to the value you were testing for. Hence, the exception, as item is 0.

    Use hasattr instead of getattr, as you want a boolean result. And for your __getitem__, you want to make sure you turn AttributeError exceptions into KeyError exceptions, to keep the interface consistent:

    from __future__ import annotations
    from dataclasses import dataclass
    
    @dataclass
    class User:
        email: int
        password: str
        blocked_at: float = None
    
        def __getitem__(self, item: str) -> str | int | float | None:
            try:
                return getattr(self, item)
            except AttributeError:
                raise KeyError(item) from None
        
        def __contains__(self, item: str) -> bool:
            return hasattr(self, item)
    

    Demo:

    >>> user = User('email@test.com', 'password')
    >>> print(user['email'])
    email@test.com
    >>> 'email' in user
    True
    >>> user["nonesuch"]
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
      File "<string>", line 12, in __getitem__
    KeyError: 'nonesuch'
    

    See the Python reference section on membership test operations for the details on how in will fall back to iteration:

    For user-defined classes which do not define __contains__() but do define __iter__(), x in y is True if some value z, for which the expression x is z or x == z is true, is produced while iterating over y. If an exception is raised during the iteration, it is as if in raised that exception.

    Lastly, the old-style iteration protocol is tried: if a class defines __getitem__(), x in y is True if and only if there is a non-negative integer index i such that x is y[i] or x == y[i], and no lower integer index raises the IndexError exception. (If any other exception is raised, it is as if in raised that exception).