Refactoring multiple IF statements to something cleaner in Python

Question

Within our codebase, we currently have a bit of code that looks like this. (Python obviously). It's now stretching to some 100 lines long, and feels very much like it is a candidate for refactoring.

TOOL_NAME = "tool_name" TOOL_NAME_TWO = "tool_name_two"

class DetailRetrievalFactory: def get_detailed_search_tool(self, tool_name):

    if tool_name == TOOL_NAME:
        return ToolNameDetailRetrieval()
    if tool_name == TOOL_NAME_TWO:
        return ToolNameTwoDetailRetrieval()
    

    raise NotImplementedError

To give some details on how its used, each subclass has its own start_detailed_search method. Looking for ideas here on how best to refactor the above, or what patterns would potentially be worth looking at to remove the growing if statement within the detail retrieval factory. A switch statement doesn't seem much better.. I was thinking is this a strategy Pattern candidate? Any advice appreciated.

data_retrieval = DetailRetrievalFactory().get_detailed_search_tool(tool_name) data_retrieval.start_detailed_search(query_id, request_id)

Answer 1

If you have many functions you can store them in a dictionary.

retrieval_funcs = {
    TOOL_NAME: ToolNameDetailRetrieval,
    TOOL_NAME_TWO: ToolNameTwoDetailRetrieval,
    TOOL_NAME_THREE: ToolNameThreeDetailRetrieval,
    TOOL_NAME_FOUR: ToolNameFourDetailRetrieval,
    TOOL_NAME_FIVE: ToolNameFiveDetailRetrieval,
    TOOL_NAME_SIX: ToolNameSixDetailRetrieval,
}

if tool_name in retrieval_funcs:
    return retrieval_funcs[tool_name]()

raise NotImplementedError