Loop to streamline pandas dataframe to_sql

Question

This code gives me what I am looking for. But I'm just thinking how I can streamline the if statements because I would be repeating myself a couple of times, and that's not really good isn't it?

import requests
import pandas
from sqlalchemy import create_engine
import os
import numpy
from selenium import webdriver
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions
from selenium.webdriver.common.by import By

def _format_data_frame(dataframe_source, df_id):
"""Format the dataframe source that is retrieved from 'my_function()' and put into SQL"""

    for each_dataframe in dataframe_source: # Formatting the columns
        if "Unnamed: 2" in each_dataframe.columns:
            each_dataframe.drop(each_dataframe.index[0], inplace=True)
            each_dataframe.rename(columns={"Fare Per Ride (cent)": "Card", "Unnamed: 2": "Cash"}, inplace=True)
        if "Card Fare Per Ride (cent)" in each_dataframe.columns:
            each_dataframe.rename(columns={"Card Fare Per Ride (cent)": "Card"}, inplace=True)
        if "Card Fare (cent)" in each_dataframe.columns:
            each_dataframe.rename(columns={"Card Fare (cent)": "Card"}, inplace=True)
        if "Description" in each_dataframe.columns:
            each_dataframe.rename(columns={"Description": "Distance"}, inplace=True)

    # Each dataframe_source has a total of 5 dataframes extracted. 
    # I don't need the last dataframe, and this portion is just to separate the dataframes out.
    truck_services = dataframe_source[0]
    feeder_services = dataframe_source[1]
    express_services = dataframe_source[2]
    other_services = dataframe_source[3]


    ### How can I streamline the below code?  ###

    if df_id == "df1":
        engine = create_engine("sqlite:///abc.db", echo=False)
        connection = engine.connect()
        pandas.DataFrame.to_sql(truck_services, name="Truck Services", con=engine, if_exists="append")
        pandas.DataFrame.to_sql(feeder_services, name="Feeder Services", con=engine, if_exists="append")
        pandas.DataFrame.to_sql(express_services, name="Express Services", con=engine, if_exists="append")
        pandas.DataFrame.to_sql(other_services, name="Other Services", con=engine, if_exists="append")
        connection.close()

    if df_id == "df2":
        engine = create_engine("sqlite:///defg_Fares.db", echo=False)
        connection = engine.connect()
        pandas.DataFrame.to_sql(truck_services, name="Truck Services", con=engine, if_exists="append")
        pandas.DataFrame.to_sql(feeder_services, name="Feeder Services", con=engine, if_exists="append")
        pandas.DataFrame.to_sql(express_services, name="Express Services", con=engine, if_exists="append")
        pandas.DataFrame.to_sql(other_services, name="Other Services", con=engine, if_exists="append")
        connection.close()

    if df_id == "df3":
        engine = create_engine("sqlite:///hijk_Fares.db", echo=False)
        ## Same thing
        connection = engine.connect()

    if df_id == "df4":
        engine = create_engine("sqlite:///lmno_Fares.db", echo=False)
        ## Same thing
        connection = engine.connect()

    if df_id == "df5":
        engine = create_engine("sqlite:///pqr_Fares.db", echo=False)
        ## Same thing
        connection = engine.connect()

Answer 1

As far as I can see the only thing changing is the URL to the DB file.

You can just define a dictionary for this:

DB_URL = {"df1": "sqlite:///abc_Fares.db",
          "df2": "sqlite:///defg_Fares.db",
          ...}

I would also define the titles of the dataframes as a list:

titles = ["Truck Services", ...]

Which you can the easily use:

engine = create_engine(DB_URL[df_id])
connection = engine.connect()
for df, title in zip(dataframe_source, titles):
    df.to_sql(title, engine, if_exists="append")
connection.close()

Note that this calls to_sql directly on the dataframes, so no need for pandas.DataFrame(df, ...).

Also note that zip will stop after the shorter iterable is exhausted. So if the list of titles only contains four titles, the fifth dataframe will not be written to the DB.

Final code:

DB_URL = {"df1": "sqlite:///abc_Fares.db",
          "df2": "sqlite:///defg_Fares.db",
          ...}

def _format_data_frame(dataframe_source, df_id):
    """Format the dataframe source that is retrieved from 'my_function()' and put into SQL"""
    column_rename = {"Fare Per Ride (cent)": "Card",
                     "Unnamed: 2": "Cash",
                     "Card Fare Per Ride (cent)": "Card",
                     "Card Fare (cent)": "Card",
                     "Description": "Distance"}

    titles = ["Truck Services", "Feeder Services", "Express Services", "Other Services"]
    engine = create_engine(DB_URL[df_id])
    connection = engine.connect()
    # Since zip stops after the shorter iterable is exhausted, this
    # leaves the fifth df out
    for df, title in zip(dataframe_source, titles):
        if "Unnamed: 2" in df.columns:
            df.drop("Unnamed: 2", axis=1, inplace=True)
        df.rename(columns=column_rename, inplace=True)
        df.to_sql(title, engine, if_exists="append")
    connection.close()

Note that I also made the if "xxx" in each_dataframe.columns faster by storing the columns once per dataframe in a set, for which membership testing is \$\mathcal{O}(1)\$.

I also made the column renaming a lot easier. The dataframe will ignore all keys in the translation dictionary for which no columns exist, so we can use one common dictionary.

You should also include your imports, right now it is not clear where create_engine comes from. You should have a look if they implemented contextmanagers, so you could do:

with engine.connect() as connection:
    for df, title in zip(dataframe_source, titles):
            df.to_sql(title, engine, if_exists="append")

Where the connection.close() is done automatically.