2
\$\begingroup\$

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()
\$\endgroup\$
0
2
\$\begingroup\$

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.

\$\endgroup\$
3
  • \$\begingroup\$ Thanks for your input! I've included my imports as you have mentioned. Am using Selenium for the create_engine. Trying out your recommendation as I type. :) \$\endgroup\$
    – jake wong
    Jan 26 '17 at 11:22
  • \$\begingroup\$ @jakewong I just updated the answer to simplify the column renaming. Also, you should be careful with updating the question after answer(s) have been posted, since you could invalidate them. \$\endgroup\$
    – Graipher
    Jan 26 '17 at 11:50
  • \$\begingroup\$ Duly noted. The last edit I have made to my question was to include the importsthat I have for the file. \$\endgroup\$
    – jake wong
    Jan 26 '17 at 11:58

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.