jojonas / pyedid Goto Github PK

import csv
import datetime
import os.path
import shutil
import sys
from argparse import ArgumentParser
from contextlib import suppress
from pathlib import Path

from bs4 import BeautifulSoup

import requests


def log(*args, **kwargs):
    return print(*args, file=sys.stderr, **kwargs)


def main():
    # Parse command line arguments
    log('Parsing command line arguments')
    parser = ArgumentParser()
    parser.add_argument('-c', '--clear-cache', action='store_true',
                        default=False)
    parser.add_argument('-a', '--include-approved', action='store_true',
                        default=False)
    parser.add_argument('-u', '--url',
                        default='https://www.uefi.org/pnp_id_list')
    parser.add_argument('cache', nargs='?', default='~/.cache/edid')
    args = parser.parse_args()
    args.cache = Path(os.path.expanduser(args.cache)).absolute()

    # Prepare cache
    if args.clear_cache:
        log(f'Clearing cache {args.cache}')
        shutil.rmtree(args.cache)
    log(f'Ensuring cache {args.cache}')
    args.cache.mkdir(parents=True, exist_ok=True)

    # Load HTML from cache or the web
    html_path = args.cache / 'src.html'
    if html_path.exists():
        log(f'Loading cached HTML from {html_path}')
        html = html_path.read_text()
    else:
        log(f'Downloading HTML from {args.url}')
        response = requests.get(args.url)
        html = response.text
        with html_path.open('w') as html_file:
            html_file.write(html)

    # Extract data from HTML
    log(f'Extracting data from HTML')
    soup = BeautifulSoup(html, 'html5lib')
    rows = []
    for tr in soup.find('tbody').find_all('tr'):
        name_td, pnp_id_td, approved_td = tr.find_all('td')
        name = name_td.text.strip()
        pnp_id = pnp_id_td.text.strip()
        row = [pnp_id, name]
        if args.include_approved:
            approved = datetime.datetime.strptime(approved_td.text.strip(),
                                                  '%m/%d/%Y').date()
            row.append(approved)
        rows.append(row)

    # Write CSV file
    log(f'Writing CSV file')
    rows.sort()
    csv.writer(sys.stdout).writerows(rows)
    sys.stdout.close()


if __name__ == '__main__':
    with suppress(KeyboardInterrupt):
        main()