Code snippet code sugestion

suggest_codes.py

from typing import Union

import pandas as pd

# Load dataset with codes order by probability
URL = ('https://gist.githubusercontent.com/pablomm/'
       '99cb95ddf65a1001a3c98312418254a0/raw/'
       '259deec38e3e1e3da72d2f6b8de8c0690b4d28f7/'
       'gistfile1.txt')

df_search = pd.read_csv(URL, index_col='index')

def suggest_codes(cod_1: Union[int, str], cod_2: Union[int, str]) -> pd.DataFrame:
    """
    Args:
        cod_1 (int): First verification mobile code.
        cod_2 (int): First verification mobile code.
    
    Returns:
        Dataframe with pin code suggestions sorted by probability.
    """
    cod_1, cod_2 = str(cod_1).zfill(4)[2:], str(cod_2).zfill(4)[2:]
    if cod_2 < cod_1:
        cod_1, cod_2 = cod_2, cod_1
    key = int(cod_1 + cod_2)

    return df_search.loc[key]