Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Python: Is is possible to have a class method that acts on a slice of the class?

Tags:

python

list

slice

I'm making a container class in Python which will either inherit from list or just implement all the standard list methods (I don't really care which).

How would I create a method that will act on only the items that are returned from a slice? I've been able to make a method that will act on the entire container (see below), but I can't seem to figure out how to act on just the slice.

I'm using python 2.7.6 and use from __future__ import print_function, division in all my code.

Example code:

from __future__ import print_function, division
import itertools

class MyContainerClass(list):
    """ For now, I'm inheriting from list. Is this my problem? """
    def __init__(self, data_array):
        list.__init__(self, data_array)

    def __getitem__(self, *args):

        arg = args[0]

        # specific indices   MyContainerClass[0, 3, 5] => indexes 0, 3, and 5
        if isinstance(arg, (list, tuple)) and not isinstance(arg[0], bool):
            return [list.__getitem__(self, _i) for _i in arg]

        # standard slice notation
        elif isinstance(arg, slice):
            return list.__getitem__(self, arg)

        # Using an array mask   MyContainerClass[[True, False, False,...,True]] => index 0 and -1]
        elif isinstance(arg[0], bool):  # or isinstance(arg, np.ndarray):
            return list(itertools.compress(self, arg))

        else:
            # I'll eventually replace this with and exception raise.
            return 'error'

    def my_method(self):
        """
        Will act on entire list, but I want it to act on only what's
        returned by the slice (which may be the entire list in some cases).
        """
        return "a, ".join([str(_i) for _i in self])

Here is an example of the usage that I'd like:

>>> data = MyContainerClass([1, 2, 3, 4, 5, 6, 7])
>>> data[5:]
[6, 7]
>>> data.my_method()           # This works as expected
"1a, 2a, 3a, 4a, 5a, 6a, 7"
>>> data[0:3].my_method()      # Doesn't work
"1a, 2a, 3"                    # but should return this


Looks like things are working now. Thanks a bunch guys! Here's what I've come up with:

from __future__ import print_function, division
import itertools

class MyContainerClass(list):
    """
    """
    def __init__(self, array):
        if isinstance(array, int):
            list.__init__(self, [array])
        else:
            list.__init__(self, array)

    def __getitem__(self, arg):
        # Standard Slice notation
        if isinstance(arg, slice):
            retval = super(MyContainerClass, self).__getitem__(arg)

        # specific indices
        elif isinstance(arg, (list, tuple)) and not isinstance(arg[0], bool):
            retval = [list.__getitem__(self, _i) for _i in arg]

        # a single specific index
        elif isinstance(arg, int):
            retval = list.__getitem__(self, arg)

        # an array mask of T/F values
        elif isinstance(arg[0], bool):      # or isinstance(arg, np.ndarray):
            retval = list(itertools.compress(self, arg))

        # raise an error on unknown
        else:
            raise SyntaxError("Unknown notation for list slice or index")

        retval = type(self)(retval)
        return retval

    def __getslice__(self, i, j):
        # Python 2 built-in types only
        return self.__getitem__(slice(i, j))

    def my_method(self):
        return "a, ".join([str(_i) for _i in self])

And acts like:

>>> data = MyContainerClass([1, 2, 3, 4, 5, 6, 7])
>>> mask = [True, True, False, True, False, False, False]
>>> print(data)
[1, 2, 3, 4, 5, 6, 7]
>>> print(type(data[5:]))
<class '__main__.MyContainerClass'>
>>> print(data.my_method())
1a, 2a, 3a, 4a, 5a, 6a, 7
>>> print(data[0:5].my_method())
1a, 2a, 3a, 4a, 5
>>> print(data[1, 5, 2].my_method())
2a, 6a, 3
>>> print(data[mask].my_method())
1a, 2a, 4
>>> print(data[2].my_method())
3
like image 549
dthor Avatar asked Feb 12 '23 04:02

dthor


1 Answers

You'll have to make sure your __getitem__ returns your type again:

class MyContainerClass(list):
    """ For now, I'm inheriting from list. Is this my problem? """
    def __init__(self, data_array):
        list.__init__(self, data_array)

    def my_method(self):
        """
        Will act on entire list, but I want it to act on only what's
        returned by the slice (which may be the entire list in some cases).
        """
        return "a, ".join([str(_i) for _i in self])

    def __getitem__(self, index):
        retval = super(MyContainerClass, self).__getitem__(index)
        if isinstance(index, slice):
            retval = type(self)(retval)
        return retval

    def __getslice__(self, i, j):
        # Python 2 built-in types only
        return self.__getitem__(slice(i, j))

The additional __getslice__ method is only needed in Python 2, and then only if you are inheriting from a type that already implements __getslice__. list is such a type.

Demo:

>>> data = MyContainerClass([1, 2, 3, 4, 5, 6, 7])
>>> data[:5]
[1, 2, 3, 4, 5]
>>> type(data[:5])
<class '__main__.MyContainerClass'>
>>> data.my_method()
'1a, 2a, 3a, 4a, 5a, 6a, 7'
>>> data[:3].my_method()
'1a, 2a, 3'
like image 116
Martijn Pieters Avatar answered Feb 15 '23 10:02

Martijn Pieters