Zaya
Zaya

Reputation: 336

How to define "in" for a python class

If I had a list of some python class that had a string field, is there a way to compare a string to a list of objects of that class using in? I was curious if there's a way to do something like this:

class Foo:
    def __init__(self, s):
        self.s = s

bar = Foo('test')
ls = [bar]

if 'test' in ls:
    print("Yay!")

by modifying the __eq__ method, or maybe even there's an __in__ method I'm unaware of

Upvotes: 4

Views: 2279

Answers (2)

Kasravnd
Kasravnd

Reputation: 107287

Your guess about modifying the __eq__ is correct. It's definitely one way to do this. You just need to check the value of __eq__'s argument with the respective attribute of your object. So as a very simple way you can implement something like following:

In [1]: class Foo:
   ...:     def __init__(self, s):
   ...:         self.s = s
   ...:
   ...:     def __eq__(self, value):
   ...:         return self.s == value
   ...:     

In [2]: bar = Foo('test')
   ...: ls = [bar]
   ...: 

In [3]: 'test' in ls
Out[3]: True

Note that this __eq__ method doesn't have any restrictions like type checking or some other error handling. If you think that they're needed in your case you can consider using try-excep and or other condition checking as well.

Demo:

In [9]: foo = Foo('test')

In [10]: bar = Foo('test')

In [11]: ls = ['test']

In [12]: bar in ls
Out[12]: True

In [13]: ls = [bar]

In [14]: foo in ls
Out[14]: True

Upvotes: 2

Samuel Dion-Girardeau
Samuel Dion-Girardeau

Reputation: 3170

You intuition about defining a custom __eq__ was right. You can achieve the following, which seems to be your objective here:

>>> class Foo:
...     def __init__(self, s):
...         self.s = s
...     def __eq__(self, other):
...         if isinstance(other, Foo):
...             return self.s == other.s
...         else:
...             return self.s == other
...
>>>
>>> bar = Foo('bar')
>>> l = [bar]
>>> bar in l
True
>>> 'bar' in l
True
>>> Foo('baz') in l
False
>>> 'baz' in l
False

I also want to point out that there is the __contains__ method, which allow you to define the membership operator's behavior (in). However this is a method on the container type, so here defining this method for your class Foo wouldn't do anything, since the in is applied to the list, and not its individual items.

Upvotes: 3

Related Questions