Decorate a read function that changes the way of opening a file

Question

I have the following function

def read_tower(path_to_file):

    with open(path_to_file,'r') as f:

        for text in f:

            text_line = text.strip().split(" ")

My idea is to use a decorator to change that function so I can read files from a zip folder like

with zipfile.ZipFile(Path(zip_filename)) as z:
    for filename in z.namelist():
        if re.search(r'.*\.mlm',filename):
            #with z.open(filename,mode="r") as f:
            data = read_tower(f)

So basically I want to decorate read_tower so I can use z.open instead of open.

I am very new to decorators yet. Thanks.

Answer 1

To do this using decorator, you may inject the ContextManager of any file type into the reading, writing (or any other processing you need) logic, something like this:

def reader(func):
    def wrap(*args, **kwargs):
        with func(*args, **kwargs) as f:
            for text in f:
                text_line = text.strip().split(" ")
            ...      
    return wrap

@reader
def file_reader(*args, **kwargs):
    return open(*args, **kwargs)

@reader
def zipfile_reader(z, *args, **kwargs):
   return z.open(*args, **kwargs)

file_content = file_reader(<path>, <mode>)
zipfile_content = zipfile_reader(zipfile_obj, <filename>, <extra_args>)

But I recommend simply to extract the common logic from the read_tower into a separate function:

def read_file_content(file):
    for text in f:
        text_line = text.strip().split(" ")
    ...

file usage:

with open('..path', mode) as f:
    read_file_content(f)

zipfile usage:

with zipfile.ZipFile(Path(zip_filename)) as z:
    for filename in z.namelist():
        if re.search(r'.*\.mlm',filename):
            with z.open(filename,mode="r") as f:
                read_file_content(f)

Note: code here is just to give idea, it is not completely working code.