pythonrefactoringpython-typing

How to refactor similar functions in Python?


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):
        ?

Solution

  • 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

    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.