How to refactor similar functions in Python?

Question

I have defined multiple functions where I search for a specific pydantic model in a list of pydantic models based on some attribute value.

SocketIOUserSessionID = str
RoomWithIndex = tuple[Room, RoomIndex]
RSStateWithIndex = tuple[RSState, int]
RSPacketSavedRecordWithIndex = tuple[RSPacketSavedRecordsContainer, int]


def find_room_by_id(self, id: UUID | str, where: list[Room]) -> RoomWithIndex | None:
    room = next(filter(lambda room: room.id == id, where), None)
    if room is None:
        return None
    index = where.index(room)
    return room, index

def find_room_by_session(self, session: SocketIOUserSessionID, where: list[Room]) -> RoomWithIndex | None:
    room = next(filter(lambda room: session in room.sessions, where), None)
    if room is None:
        return None
    index = where.index(room)
    return room, index

def find_rs_state_by_room_id(self, room_id: str, where: list[RSState]) -> RSStateWithIndex | None:
    rs_state = next(filter(lambda rs_state: rs_state.room_id == room_id, where), None)
    if rs_state is None:
        return None
    index = where.index(rs_state)
    return rs_state, index

def find_saved_record_by_room_id(self, room_id: str, where: list[RSPacketSavedRecordsContainer]) -> RSPacketSavedRecordWithIndex | None:
    saved_record = next(filter(lambda saved_records: saved_records.room_id == room_id, where), None)
    if saved_record is None:
        return None
    index = where.index(saved_record)
    return saved_record, index

How to write a generic function (with typing) to refactor such code? I heard of functools.singledispatch decorator but I am not sure that this is the right case to use it.

    def find_value_by_attr(self):
        ?

Answer 1

I tried to generalize the four functions as much as possible - here's where I ended up:

def find_model(
    self,
    iid: UUID | str,
    where: list[Any],
    filter_attr: str,
) -> tuple[Any, int] | None:
    if (
        model := next(
            filter(lambda model: iid == getattr(model, filter_attr), where)
        )
    ):
        return model, where.index(model)

The find_model function takes

• an iid (item ID) in place of id, session, or room_id (NOTE: naming this iid avoids shadowing the builtin id)
• a where list just as before
• a filter_attr which is the string name of the attribute you want to filter, e.g. 'id', 'sessions', or 'room_id'

The only thing this doesn't quite cover is the filtering case of find_room_by_session which is using in instead of == in the filter lambda. If anyone more clever than I would care to weigh in, I'm open to it!

If I may take a moment to editorialize: I like the idea of using @singledispatch for this, but it doesn't get you away from writing multiple function prototypes anyway...if the goal it 'less code', then @singledispatch doesn't help much in that regard.

Answer 2

To me it seems that the only way of doing it is to create another function that takes in a conditional function:

def find_value_by_attr(self, where, f):
    val = next(filter(f, where), None)
    if val is None:
        return None
    index = where.index(val)
    return record, index


def find_room_by_id(self, id: UUID | str, where: list[Room]) -> RoomWithIndex | None:
    f = lambda room: room.id == id
    return find_value_by_attr(where, f)

def find_room_by_session(self, session: SocketIOUserSessionID, where: list[Room]) -> RoomWithIndex | None:
    f = lambda room: session in room.sessions
    return find_value_by_attr(where, f)
...

Because in case of just one-for-all function you will need to join your typings which will lead to a situation when they can mix like for example UUID | str for id in the first func and RSStateWithIndex for the output of the third one. It'll break your logic.