pyjg pyjg - 1 year ago 61
Python Question

Ensuring setup function runs every time imported module function is called

I am importing a module with a few functions which connect to a database. I call functions from that module like so:

main.py:


import data_source as DS

DS.get_data(type_of_data)
DS.get_some_other_data(type_of_data)


data_source.py:


def get_data(type_of_data):
# connect to db and return a dictionary

def get_some_other_data(type_of_data):
# connect to db and return a dictionary


Every function established a database connection. I am trying to reduce code duplication so I can have a function which connects to the DB like so:

in
data_source.py
:

def connect_to_db():
# connect to bd and return connection to caller


but i would have to call
connect_to_db()
from every
function
. Is there a way to ensure if a function from an imported module is called another function runs by default? so I dont have to have
connect_to_db
as the top like in every function within
data_source.py
? and so the connection isn't made for every single call?

I know it wouldnt hurt the performance for my purposes but I am just curious. I know I can pass the connection back to the caller after the first call and then for the subsequent call(s) I can pass the connection to the function in the imported module but that is something I want to avoid doing.

like when you do
python main.py
everything under
if __name__ == "main":
is run is there perhaps something like

if __function_is_being_called__:
?

Answer Source

You can use a decorator to first call connect_to_db, then call the function as normal:

import functools

def with_db_connection(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        connect_to_db()
        return f(*args, **kwargs)
    return wrapper

@with_db_connection
def get_data(type_of_data):
    ...

You can also make it pass a database argument to avoid global variables:

import functools

def with_db_connection(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        db = connect_to_db()
        return f(db, *args, **kwargs)
    return wrapper

@with_db_connection
def get_data(db, type_of_data):
    ...

get_data('string')  # db argument is passed in through the decorator
Recommended from our users: Dynamic Network Monitoring from WhatsUp Gold from IPSwitch. Free Download