Extracting the src attribute and content of <script> tags in HTML files in a directory

Question

I have a program which parse HTML content to get a script src attribute and script attribute. A function have directory (contain HTML files) as input and extract to other directory as output.

However, I think my code is not clean. The following is my code.

def extract_int_js(src, dst):

    for file in os.listdir(src):
        filepath = os.path.join(src, file)
        with open(filepath) as f:
            html = f.read()
        parser = HtmlParser(html)

        for text in parser.script_text:
            yield text


def extract_src_js(src, dst):
    for file in os.listdir(src):
        filepath = os.path.join(src, file)
        with open(filepath) as f:
            html = f.read()
        parser = HtmlParser(html)

        for src in parser.src:
            yield src

Do you have any solution to improve this code since there are duplicated line in 2 functions.

Answer 1

You can use getattr function to merge your two functions into a single one.

def extract_attribute(src, attr_name):
    for file in os.listdir(src):
        filepath = os.path.join(src, file)
    with open(filepath) as f:
        html = f.read()
    parser = HtmlParser(html)

    yield from getattr(parser, attr_name)

Note that in Python 3.3+ there is yield from statement so you can yield values directly from iterable without writing for loop

Now if you still want to have two separate functions instead of passing attribute_name argument each time you can use function partial from functools

from functools import partial

extract_src_js = partial(extract_attribute, attr_name='src')
extract_int_js = partial(extract_attribute, attr_name='script_text')

Please also note that I've removed dst parameter from your functions since it was not used, however if it's just a part of your function you wanted us to review, just return it back.

Answer 2

If you always need both attributes and not one individually, and both have the same length, you could just iterate over both at the same time:

def extract_int_js(src, dst):

    for file in os.listdir(src):
        filepath = os.path.join(src, file)
        with open(filepath) as f:
            html = f.read()
        parser = HtmlParser(html)

        for src, text in zip(parser.script_text, parser.src):
            yield src, text

You can use yield from here as well and factor out the reading of the files to another generator:

def read_files(src, dst):
    for file in os.listdir(src):
        filepath = os.path.join(src, file)
        with open(filepath) as f:
            yield f.read()


def extract_int_js(src, dst):
    for html in read_files(src, dst):
        parser = HtmlParser(html)
        yield from zip(parser.script_text, parser.src)

These can be used like this:

for src, text in extract_int_js(src, dst):
    print(src, text)