mirror of
https://github.com/ArchiveBox/ArchiveBox
synced 2024-11-22 20:23:12 +00:00
add mypy type hints
This commit is contained in:
parent
58c9b47d43
commit
63abc8980c
3 changed files with 171 additions and 143 deletions
|
@ -1,5 +1,7 @@
|
|||
import os
|
||||
import json
|
||||
|
||||
from typing import Union, Dict, List, Tuple, NamedTuple
|
||||
from collections import defaultdict
|
||||
from datetime import datetime
|
||||
|
||||
|
@ -40,13 +42,15 @@ from util import (
|
|||
without_query,
|
||||
without_fragment,
|
||||
fetch_page_title,
|
||||
read_js_script,
|
||||
is_static_file,
|
||||
TimedProgress,
|
||||
chmod_file,
|
||||
wget_output_path,
|
||||
chrome_args,
|
||||
check_link_structure,
|
||||
run, PIPE, DEVNULL
|
||||
run, PIPE, DEVNULL,
|
||||
Link,
|
||||
)
|
||||
from logs import (
|
||||
log_link_archiving_started,
|
||||
|
@ -55,15 +59,22 @@ from logs import (
|
|||
log_archive_method_finished,
|
||||
)
|
||||
|
||||
|
||||
|
||||
class ArchiveError(Exception):
|
||||
def __init__(self, message, hints=None):
|
||||
super().__init__(message)
|
||||
self.hints = hints
|
||||
|
||||
class ArchiveResult(NamedTuple):
|
||||
cmd: List[str]
|
||||
pwd: str
|
||||
output: Union[str, Exception, None]
|
||||
status: str
|
||||
start_ts: datetime
|
||||
end_ts: datetime
|
||||
duration: int
|
||||
|
||||
def archive_link(link_dir, link):
|
||||
|
||||
def archive_link(link_dir: str, link: Link, page=None) -> Link:
|
||||
"""download the DOM, PDF, and a screenshot into a folder named after the link's timestamp"""
|
||||
|
||||
ARCHIVE_METHODS = (
|
||||
|
@ -95,10 +106,11 @@ def archive_link(link_dir, link):
|
|||
log_archive_method_started(method_name)
|
||||
|
||||
result = method_function(link_dir, link)
|
||||
link['history'][method_name].append(result)
|
||||
|
||||
stats[result['status']] += 1
|
||||
log_archive_method_finished(result)
|
||||
link['history'][method_name].append(result._asdict())
|
||||
|
||||
stats[result.status] += 1
|
||||
log_archive_method_finished(result._asdict())
|
||||
else:
|
||||
stats['skipped'] += 1
|
||||
|
||||
|
@ -117,7 +129,7 @@ def archive_link(link_dir, link):
|
|||
|
||||
### Archive Method Functions
|
||||
|
||||
def should_fetch_title(link_dir, link):
|
||||
def should_fetch_title(link_dir: str, link: Link) -> bool:
|
||||
# if link already has valid title, skip it
|
||||
if link['title'] and not link['title'].lower().startswith('http'):
|
||||
return False
|
||||
|
@ -127,7 +139,7 @@ def should_fetch_title(link_dir, link):
|
|||
|
||||
return FETCH_TITLE
|
||||
|
||||
def fetch_title(link_dir, link, timeout=TIMEOUT):
|
||||
def fetch_title(link_dir: str, link: Link, timeout: int=TIMEOUT) -> ArchiveResult:
|
||||
"""try to guess the page's title from its content"""
|
||||
|
||||
output = None
|
||||
|
@ -150,22 +162,22 @@ def fetch_title(link_dir, link, timeout=TIMEOUT):
|
|||
finally:
|
||||
timer.end()
|
||||
|
||||
return {
|
||||
'cmd': cmd,
|
||||
'pwd': link_dir,
|
||||
'output': output,
|
||||
'status': status,
|
||||
return ArchiveResult(
|
||||
cmd=cmd,
|
||||
pwd=link_dir,
|
||||
output=output,
|
||||
status=status,
|
||||
**timer.stats,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def should_fetch_favicon(link_dir, link):
|
||||
def should_fetch_favicon(link_dir: str, link: Link) -> bool:
|
||||
if os.path.exists(os.path.join(link_dir, 'favicon.ico')):
|
||||
return False
|
||||
|
||||
return FETCH_FAVICON
|
||||
|
||||
def fetch_favicon(link_dir, link, timeout=TIMEOUT):
|
||||
def fetch_favicon(link_dir: str, link: Link, timeout: int=TIMEOUT) -> ArchiveResult:
|
||||
"""download site favicon from google's favicon api"""
|
||||
|
||||
output = 'favicon.ico'
|
||||
|
@ -188,15 +200,15 @@ def fetch_favicon(link_dir, link, timeout=TIMEOUT):
|
|||
finally:
|
||||
timer.end()
|
||||
|
||||
return {
|
||||
'cmd': cmd,
|
||||
'pwd': link_dir,
|
||||
'output': output,
|
||||
'status': status,
|
||||
return ArchiveResult(
|
||||
cmd=cmd,
|
||||
pwd=link_dir,
|
||||
output=output,
|
||||
status=status,
|
||||
**timer.stats,
|
||||
}
|
||||
)
|
||||
|
||||
def should_fetch_wget(link_dir, link):
|
||||
def should_fetch_wget(link_dir: str, link: Link) -> bool:
|
||||
output_path = wget_output_path(link)
|
||||
if output_path and os.path.exists(os.path.join(link_dir, output_path)):
|
||||
return False
|
||||
|
@ -204,7 +216,7 @@ def should_fetch_wget(link_dir, link):
|
|||
return FETCH_WGET
|
||||
|
||||
|
||||
def fetch_wget(link_dir, link, timeout=TIMEOUT):
|
||||
def fetch_wget(link_dir: str, link: Link, timeout: int=TIMEOUT) -> ArchiveResult:
|
||||
"""download full site using wget"""
|
||||
|
||||
if FETCH_WARC:
|
||||
|
@ -274,15 +286,15 @@ def fetch_wget(link_dir, link, timeout=TIMEOUT):
|
|||
finally:
|
||||
timer.end()
|
||||
|
||||
return {
|
||||
'cmd': cmd,
|
||||
'pwd': link_dir,
|
||||
'output': output,
|
||||
'status': status,
|
||||
return ArchiveResult(
|
||||
cmd=cmd,
|
||||
pwd=link_dir,
|
||||
output=output,
|
||||
status=status,
|
||||
**timer.stats,
|
||||
}
|
||||
)
|
||||
|
||||
def should_fetch_pdf(link_dir, link):
|
||||
def should_fetch_pdf(link_dir: str, link: Link) -> bool:
|
||||
if is_static_file(link['url']):
|
||||
return False
|
||||
|
||||
|
@ -292,7 +304,7 @@ def should_fetch_pdf(link_dir, link):
|
|||
return FETCH_PDF
|
||||
|
||||
|
||||
def fetch_pdf(link_dir, link, timeout=TIMEOUT):
|
||||
def fetch_pdf(link_dir: str, link: Link, timeout: int=TIMEOUT) -> ArchiveResult:
|
||||
"""print PDF of site to file using chrome --headless"""
|
||||
|
||||
output = 'output.pdf'
|
||||
|
@ -317,15 +329,15 @@ def fetch_pdf(link_dir, link, timeout=TIMEOUT):
|
|||
finally:
|
||||
timer.end()
|
||||
|
||||
return {
|
||||
'cmd': cmd,
|
||||
'pwd': link_dir,
|
||||
'output': output,
|
||||
'status': status,
|
||||
return ArchiveResult(
|
||||
cmd=cmd,
|
||||
pwd=link_dir,
|
||||
output=output,
|
||||
status=status,
|
||||
**timer.stats,
|
||||
}
|
||||
)
|
||||
|
||||
def should_fetch_screenshot(link_dir, link):
|
||||
def should_fetch_screenshot(link_dir: str, link: Link) -> bool:
|
||||
if is_static_file(link['url']):
|
||||
return False
|
||||
|
||||
|
@ -334,7 +346,7 @@ def should_fetch_screenshot(link_dir, link):
|
|||
|
||||
return FETCH_SCREENSHOT
|
||||
|
||||
def fetch_screenshot(link_dir, link, timeout=TIMEOUT):
|
||||
def fetch_screenshot(link_dir: str, link: Link, timeout: int=TIMEOUT) -> ArchiveResult:
|
||||
"""take screenshot of site using chrome --headless"""
|
||||
|
||||
output = 'screenshot.png'
|
||||
|
@ -359,15 +371,15 @@ def fetch_screenshot(link_dir, link, timeout=TIMEOUT):
|
|||
finally:
|
||||
timer.end()
|
||||
|
||||
return {
|
||||
'cmd': cmd,
|
||||
'pwd': link_dir,
|
||||
'output': output,
|
||||
'status': status,
|
||||
return ArchiveResult(
|
||||
cmd=cmd,
|
||||
pwd=link_dir,
|
||||
output=output,
|
||||
status=status,
|
||||
**timer.stats,
|
||||
}
|
||||
)
|
||||
|
||||
def should_fetch_dom(link_dir, link):
|
||||
def should_fetch_dom(link_dir: str, link: Link) -> bool:
|
||||
if is_static_file(link['url']):
|
||||
return False
|
||||
|
||||
|
@ -376,7 +388,7 @@ def should_fetch_dom(link_dir, link):
|
|||
|
||||
return FETCH_DOM
|
||||
|
||||
def fetch_dom(link_dir, link, timeout=TIMEOUT):
|
||||
def fetch_dom(link_dir: str, link: Link, timeout: int=TIMEOUT) -> ArchiveResult:
|
||||
"""print HTML of site to file using chrome --dump-html"""
|
||||
|
||||
output = 'output.html'
|
||||
|
@ -403,15 +415,15 @@ def fetch_dom(link_dir, link, timeout=TIMEOUT):
|
|||
finally:
|
||||
timer.end()
|
||||
|
||||
return {
|
||||
'cmd': cmd,
|
||||
'pwd': link_dir,
|
||||
'output': output,
|
||||
'status': status,
|
||||
return ArchiveResult(
|
||||
cmd=cmd,
|
||||
pwd=link_dir,
|
||||
output=output,
|
||||
status=status,
|
||||
**timer.stats,
|
||||
}
|
||||
)
|
||||
|
||||
def should_fetch_git(link_dir, link):
|
||||
def should_fetch_git(link_dir: str, link: Link) -> bool:
|
||||
if is_static_file(link['url']):
|
||||
return False
|
||||
|
||||
|
@ -428,7 +440,7 @@ def should_fetch_git(link_dir, link):
|
|||
return FETCH_GIT
|
||||
|
||||
|
||||
def fetch_git(link_dir, link, timeout=TIMEOUT):
|
||||
def fetch_git(link_dir: str, link: Link, timeout: int=TIMEOUT) -> ArchiveResult:
|
||||
"""download full site using git"""
|
||||
|
||||
output = 'git'
|
||||
|
@ -460,16 +472,16 @@ def fetch_git(link_dir, link, timeout=TIMEOUT):
|
|||
finally:
|
||||
timer.end()
|
||||
|
||||
return {
|
||||
'cmd': cmd,
|
||||
'pwd': link_dir,
|
||||
'output': output,
|
||||
'status': status,
|
||||
return ArchiveResult(
|
||||
cmd=cmd,
|
||||
pwd=link_dir,
|
||||
output=output,
|
||||
status=status,
|
||||
**timer.stats,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def should_fetch_media(link_dir, link):
|
||||
def should_fetch_media(link_dir: str, link: Link) -> bool:
|
||||
if is_static_file(link['url']):
|
||||
return False
|
||||
|
||||
|
@ -478,7 +490,7 @@ def should_fetch_media(link_dir, link):
|
|||
|
||||
return FETCH_MEDIA
|
||||
|
||||
def fetch_media(link_dir, link, timeout=MEDIA_TIMEOUT):
|
||||
def fetch_media(link_dir: str, link: Link, timeout: int=MEDIA_TIMEOUT) -> ArchiveResult:
|
||||
"""Download playlists or individual video, audio, and subtitles using youtube-dl"""
|
||||
|
||||
output = 'media'
|
||||
|
@ -531,16 +543,16 @@ def fetch_media(link_dir, link, timeout=MEDIA_TIMEOUT):
|
|||
finally:
|
||||
timer.end()
|
||||
|
||||
return {
|
||||
'cmd': cmd,
|
||||
'pwd': link_dir,
|
||||
'output': output,
|
||||
'status': status,
|
||||
return ArchiveResult(
|
||||
cmd=cmd,
|
||||
pwd=link_dir,
|
||||
output=output,
|
||||
status=status,
|
||||
**timer.stats,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def should_fetch_archive_dot_org(link_dir, link):
|
||||
def should_fetch_archive_dot_org(link_dir: str, link: Link) -> bool:
|
||||
if is_static_file(link['url']):
|
||||
return False
|
||||
|
||||
|
@ -550,7 +562,7 @@ def should_fetch_archive_dot_org(link_dir, link):
|
|||
|
||||
return SUBMIT_ARCHIVE_DOT_ORG
|
||||
|
||||
def archive_dot_org(link_dir, link, timeout=TIMEOUT):
|
||||
def archive_dot_org(link_dir: str, link: Link, timeout: int=TIMEOUT) -> ArchiveResult:
|
||||
"""submit site to archive.org for archiving via their service, save returned archive url"""
|
||||
|
||||
output = 'archive.org.txt'
|
||||
|
@ -596,17 +608,17 @@ def archive_dot_org(link_dir, link, timeout=TIMEOUT):
|
|||
chmod_file('archive.org.txt', cwd=link_dir)
|
||||
output = archive_org_url
|
||||
|
||||
return {
|
||||
'cmd': cmd,
|
||||
'pwd': link_dir,
|
||||
'output': output,
|
||||
'status': status,
|
||||
return ArchiveResult(
|
||||
cmd=cmd,
|
||||
pwd=link_dir,
|
||||
output=output,
|
||||
status=status,
|
||||
**timer.stats,
|
||||
}
|
||||
)
|
||||
|
||||
def parse_archive_dot_org_response(response):
|
||||
def parse_archive_dot_org_response(response: bytes) -> Tuple[List[str], List[str]]:
|
||||
# Parse archive.org response headers
|
||||
headers = defaultdict(list)
|
||||
headers: Dict[str, List[str]] = defaultdict(list)
|
||||
|
||||
# lowercase all the header names and store in dict
|
||||
for header in response.splitlines():
|
||||
|
|
|
@ -3,6 +3,8 @@ import json
|
|||
|
||||
from datetime import datetime
|
||||
from string import Template
|
||||
from typing import List, Tuple
|
||||
|
||||
try:
|
||||
from distutils.dir_util import copy_tree
|
||||
except ImportError:
|
||||
|
@ -23,6 +25,7 @@ from util import (
|
|||
check_links_structure,
|
||||
wget_output_path,
|
||||
latest_output,
|
||||
Link,
|
||||
)
|
||||
from parse import parse_links
|
||||
from links import validate_links
|
||||
|
@ -39,7 +42,7 @@ TITLE_LOADING_MSG = 'Not yet archived...'
|
|||
|
||||
### Homepage index for all the links
|
||||
|
||||
def write_links_index(out_dir, links, finished=False):
|
||||
def write_links_index(out_dir: str, links: List[Link], finished: bool=False) -> None:
|
||||
"""create index.html file for a given list of links"""
|
||||
|
||||
log_indexing_process_started()
|
||||
|
@ -53,15 +56,15 @@ def write_links_index(out_dir, links, finished=False):
|
|||
write_html_links_index(out_dir, links, finished=finished)
|
||||
log_indexing_finished(out_dir, 'index.html')
|
||||
|
||||
def load_links_index(out_dir=OUTPUT_DIR, import_path=None):
|
||||
def load_links_index(out_dir: str=OUTPUT_DIR, import_path: str=None) -> Tuple[List[Link], List[Link]]:
|
||||
"""parse and load existing index with any new links from import_path merged in"""
|
||||
|
||||
existing_links = []
|
||||
existing_links: List[Link] = []
|
||||
if out_dir:
|
||||
existing_links = parse_json_links_index(out_dir)
|
||||
check_links_structure(existing_links)
|
||||
|
||||
new_links = []
|
||||
new_links: List[Link] = []
|
||||
if import_path:
|
||||
# parse and validate the import file
|
||||
log_parsing_started(import_path)
|
||||
|
@ -79,7 +82,7 @@ def load_links_index(out_dir=OUTPUT_DIR, import_path=None):
|
|||
|
||||
return all_links, new_links
|
||||
|
||||
def write_json_links_index(out_dir, links):
|
||||
def write_json_links_index(out_dir: str, links: List[Link]) -> None:
|
||||
"""write the json link index to a given path"""
|
||||
|
||||
check_links_structure(links)
|
||||
|
@ -100,7 +103,7 @@ def write_json_links_index(out_dir, links):
|
|||
|
||||
chmod_file(path)
|
||||
|
||||
def parse_json_links_index(out_dir=OUTPUT_DIR):
|
||||
def parse_json_links_index(out_dir: str=OUTPUT_DIR) -> List[Link]:
|
||||
"""parse a archive index json file and return the list of links"""
|
||||
index_path = os.path.join(out_dir, 'index.json')
|
||||
if os.path.exists(index_path):
|
||||
|
@ -111,7 +114,7 @@ def parse_json_links_index(out_dir=OUTPUT_DIR):
|
|||
|
||||
return []
|
||||
|
||||
def write_html_links_index(out_dir, links, finished=False):
|
||||
def write_html_links_index(out_dir: str, links: List[Link], finished: bool=False) -> None:
|
||||
"""write the html link index to a given path"""
|
||||
|
||||
check_links_structure(links)
|
||||
|
@ -166,7 +169,7 @@ def write_html_links_index(out_dir, links, finished=False):
|
|||
chmod_file(path)
|
||||
|
||||
|
||||
def patch_links_index(link, out_dir=OUTPUT_DIR):
|
||||
def patch_links_index(link: Link, out_dir: str=OUTPUT_DIR) -> None:
|
||||
"""hack to in-place update one row's info in the generated index html"""
|
||||
|
||||
title = link['title'] or latest_output(link)['title']
|
||||
|
@ -200,12 +203,12 @@ def patch_links_index(link, out_dir=OUTPUT_DIR):
|
|||
|
||||
### Individual link index
|
||||
|
||||
def write_link_index(out_dir, link):
|
||||
def write_link_index(out_dir: str, link: Link) -> None:
|
||||
link['updated'] = str(datetime.now().timestamp())
|
||||
write_json_link_index(out_dir, link)
|
||||
write_html_link_index(out_dir, link)
|
||||
|
||||
def write_json_link_index(out_dir, link):
|
||||
def write_json_link_index(out_dir: str, link: Link) -> None:
|
||||
"""write a json file with some info about the link"""
|
||||
|
||||
check_link_structure(link)
|
||||
|
@ -216,7 +219,7 @@ def write_json_link_index(out_dir, link):
|
|||
|
||||
chmod_file(path)
|
||||
|
||||
def parse_json_link_index(out_dir):
|
||||
def parse_json_link_index(out_dir: str) -> dict:
|
||||
"""load the json link index from a given directory"""
|
||||
existing_index = os.path.join(out_dir, 'index.json')
|
||||
if os.path.exists(existing_index):
|
||||
|
@ -226,7 +229,7 @@ def parse_json_link_index(out_dir):
|
|||
return link_json
|
||||
return {}
|
||||
|
||||
def load_json_link_index(out_dir, link):
|
||||
def load_json_link_index(out_dir: str, link: Link) -> Link:
|
||||
"""check for an existing link archive in the given directory,
|
||||
and load+merge it into the given link dict
|
||||
"""
|
||||
|
@ -241,7 +244,7 @@ def load_json_link_index(out_dir, link):
|
|||
check_link_structure(link)
|
||||
return link
|
||||
|
||||
def write_html_link_index(out_dir, link):
|
||||
def write_html_link_index(out_dir: str, link: Link) -> None:
|
||||
check_link_structure(link)
|
||||
with open(os.path.join(TEMPLATES_DIR, 'link_index.html'), 'r', encoding='utf-8') as f:
|
||||
link_html = f.read()
|
||||
|
|
|
@ -3,6 +3,8 @@ import re
|
|||
import sys
|
||||
import time
|
||||
|
||||
from typing import List, Dict, Any, Optional, Union
|
||||
|
||||
from urllib.request import Request, urlopen
|
||||
from urllib.parse import urlparse, quote
|
||||
from decimal import Decimal
|
||||
|
@ -30,6 +32,7 @@ from config import (
|
|||
CHECK_SSL_VALIDITY,
|
||||
WGET_USER_AGENT,
|
||||
CHROME_OPTIONS,
|
||||
PYTHON_PATH,
|
||||
)
|
||||
from logs import pretty_path
|
||||
|
||||
|
@ -86,9 +89,11 @@ STATICFILE_EXTENSIONS = {
|
|||
# html, htm, shtml, xhtml, xml, aspx, php, cgi
|
||||
}
|
||||
|
||||
Link = Dict[str, Any]
|
||||
|
||||
### Checks & Tests
|
||||
|
||||
def check_link_structure(link):
|
||||
def check_link_structure(link: Link) -> None:
|
||||
"""basic sanity check invariants to make sure the data is valid"""
|
||||
assert isinstance(link, dict)
|
||||
assert isinstance(link.get('url'), str)
|
||||
|
@ -100,13 +105,13 @@ def check_link_structure(link):
|
|||
assert isinstance(key, str)
|
||||
assert isinstance(val, list), 'history must be a Dict[str, List], got: {}'.format(link['history'])
|
||||
|
||||
def check_links_structure(links):
|
||||
def check_links_structure(links: List[Link]) -> None:
|
||||
"""basic sanity check invariants to make sure the data is valid"""
|
||||
assert isinstance(links, list)
|
||||
if links:
|
||||
check_link_structure(links[0])
|
||||
|
||||
def check_url_parsing_invariants():
|
||||
def check_url_parsing_invariants() -> None:
|
||||
"""Check that plain text regex URL parsing works as expected"""
|
||||
|
||||
# this is last-line-of-defense to make sure the URL_REGEX isn't
|
||||
|
@ -137,7 +142,7 @@ def check_url_parsing_invariants():
|
|||
|
||||
### Random Helpers
|
||||
|
||||
def save_stdin_source(raw_text):
|
||||
def save_stdin_source(raw_text: str) -> str:
|
||||
if not os.path.exists(SOURCES_DIR):
|
||||
os.makedirs(SOURCES_DIR)
|
||||
|
||||
|
@ -150,7 +155,7 @@ def save_stdin_source(raw_text):
|
|||
|
||||
return source_path
|
||||
|
||||
def save_remote_source(url, timeout=TIMEOUT):
|
||||
def save_remote_source(url: str, timeout: int=TIMEOUT) -> str:
|
||||
"""download a given url's content into output/sources/domain-<timestamp>.txt"""
|
||||
|
||||
if not os.path.exists(SOURCES_DIR):
|
||||
|
@ -187,7 +192,7 @@ def save_remote_source(url, timeout=TIMEOUT):
|
|||
|
||||
return source_path
|
||||
|
||||
def fetch_page_title(url, timeout=10, progress=SHOW_PROGRESS):
|
||||
def fetch_page_title(url: str, timeout: int=10, progress: bool=SHOW_PROGRESS) -> Optional[str]:
|
||||
"""Attempt to guess a page's title by downloading the html"""
|
||||
|
||||
if not FETCH_TITLE:
|
||||
|
@ -209,7 +214,7 @@ def fetch_page_title(url, timeout=10, progress=SHOW_PROGRESS):
|
|||
# ))
|
||||
return None
|
||||
|
||||
def wget_output_path(link):
|
||||
def wget_output_path(link: Link) -> Optional[str]:
|
||||
"""calculate the path to the wgetted .html file, since wget may
|
||||
adjust some paths to be different than the base_url path.
|
||||
|
||||
|
@ -278,9 +283,15 @@ def wget_output_path(link):
|
|||
return None
|
||||
|
||||
|
||||
def read_js_script(script_name: str) -> str:
|
||||
script_path = os.path.join(PYTHON_PATH, 'scripts', script_name)
|
||||
|
||||
with open(script_path, 'r') as f:
|
||||
return f.read().split('// INFO BELOW HERE')[0].strip()
|
||||
|
||||
### String Manipulation & Logging Helpers
|
||||
|
||||
def str_between(string, start, end=None):
|
||||
def str_between(string: str, start: str, end: str=None) -> str:
|
||||
"""(<abc>12345</def>, <abc>, </def>) -> 12345"""
|
||||
|
||||
content = string.split(start, 1)[-1]
|
||||
|
@ -292,7 +303,7 @@ def str_between(string, start, end=None):
|
|||
|
||||
### Link Helpers
|
||||
|
||||
def merge_links(a, b):
|
||||
def merge_links(a: Link, b: Link) -> Link:
|
||||
"""deterministially merge two links, favoring longer field values over shorter,
|
||||
and "cleaner" values over worse ones.
|
||||
"""
|
||||
|
@ -310,7 +321,7 @@ def merge_links(a, b):
|
|||
'sources': list(set(a.get('sources', []) + b.get('sources', []))),
|
||||
}
|
||||
|
||||
def is_static_file(url):
|
||||
def is_static_file(url: str) -> bool:
|
||||
"""Certain URLs just point to a single static file, and
|
||||
don't need to be re-archived in many formats
|
||||
"""
|
||||
|
@ -318,7 +329,7 @@ def is_static_file(url):
|
|||
# TODO: the proper way is with MIME type detection, not using extension
|
||||
return extension(url) in STATICFILE_EXTENSIONS
|
||||
|
||||
def derived_link_info(link):
|
||||
def derived_link_info(link: Link) -> dict:
|
||||
"""extend link info with the archive urls and other derived data"""
|
||||
|
||||
url = link['url']
|
||||
|
@ -373,7 +384,7 @@ def derived_link_info(link):
|
|||
return extended_info
|
||||
|
||||
|
||||
def latest_output(link, status=None):
|
||||
def latest_output(link: Link, status: str=None) -> Dict[str, Optional[str]]:
|
||||
"""get the latest output that each archive method produced for link"""
|
||||
|
||||
latest = {
|
||||
|
@ -440,7 +451,42 @@ def run(*popenargs, input=None, capture_output=False, timeout=None, check=False,
|
|||
return CompletedProcess(process.args, retcode, stdout, stderr)
|
||||
|
||||
|
||||
def progress_bar(seconds, prefix):
|
||||
class TimedProgress:
|
||||
"""Show a progress bar and measure elapsed time until .end() is called"""
|
||||
|
||||
def __init__(self, seconds, prefix=''):
|
||||
if SHOW_PROGRESS:
|
||||
self.p = Process(target=progress_bar, args=(seconds, prefix))
|
||||
self.p.start()
|
||||
|
||||
self.stats = {
|
||||
'start_ts': datetime.now(),
|
||||
'end_ts': None,
|
||||
'duration': None,
|
||||
}
|
||||
|
||||
def end(self):
|
||||
"""immediately end progress, clear the progressbar line, and save end_ts"""
|
||||
|
||||
end_ts = datetime.now()
|
||||
self.stats.update({
|
||||
'end_ts': end_ts,
|
||||
'duration': (end_ts - self.stats['start_ts']).seconds,
|
||||
})
|
||||
|
||||
if SHOW_PROGRESS:
|
||||
# protect from double termination
|
||||
#if p is None or not hasattr(p, 'kill'):
|
||||
# return
|
||||
if self.p is not None:
|
||||
self.p.terminate()
|
||||
self.p = None
|
||||
|
||||
sys.stdout.write('\r{}{}\r'.format((' ' * TERM_WIDTH), ANSI['reset'])) # clear whole terminal line
|
||||
sys.stdout.flush()
|
||||
|
||||
|
||||
def progress_bar(seconds: int, prefix: str='') -> None:
|
||||
"""show timer in the form of progress bar, with percentage and seconds remaining"""
|
||||
chunk = '█' if sys.stdout.encoding == 'UTF-8' else '#'
|
||||
chunks = TERM_WIDTH - len(prefix) - 20 # number of progress chunks to show (aka max bar width)
|
||||
|
@ -477,41 +523,8 @@ def progress_bar(seconds, prefix):
|
|||
print()
|
||||
pass
|
||||
|
||||
class TimedProgress:
|
||||
"""Show a progress bar and measure elapsed time until .end() is called"""
|
||||
|
||||
def __init__(self, seconds, prefix=''):
|
||||
if SHOW_PROGRESS:
|
||||
self.p = Process(target=progress_bar, args=(seconds, prefix))
|
||||
self.p.start()
|
||||
|
||||
self.stats = {
|
||||
'start_ts': datetime.now(),
|
||||
'end_ts': None,
|
||||
'duration': None,
|
||||
}
|
||||
|
||||
def end(self):
|
||||
"""immediately end progress, clear the progressbar line, and save end_ts"""
|
||||
|
||||
end_ts = datetime.now()
|
||||
self.stats.update({
|
||||
'end_ts': end_ts,
|
||||
'duration': (end_ts - self.stats['start_ts']).seconds,
|
||||
})
|
||||
|
||||
if SHOW_PROGRESS:
|
||||
# protect from double termination
|
||||
#if p is None or not hasattr(p, 'kill'):
|
||||
# return
|
||||
if self.p is not None:
|
||||
self.p.terminate()
|
||||
self.p = None
|
||||
|
||||
sys.stdout.write('\r{}{}\r'.format((' ' * TERM_WIDTH), ANSI['reset'])) # clear whole terminal line
|
||||
sys.stdout.flush()
|
||||
|
||||
def download_url(url, timeout=TIMEOUT):
|
||||
def download_url(url: str, timeout: int=TIMEOUT) -> str:
|
||||
"""Download the contents of a remote url and return the text"""
|
||||
|
||||
req = Request(url, headers={'User-Agent': WGET_USER_AGENT})
|
||||
|
@ -526,7 +539,7 @@ def download_url(url, timeout=TIMEOUT):
|
|||
encoding = resp.headers.get_content_charset() or 'utf-8'
|
||||
return resp.read().decode(encoding)
|
||||
|
||||
def chmod_file(path, cwd='.', permissions=OUTPUT_PERMISSIONS, timeout=30):
|
||||
def chmod_file(path: str, cwd: str='.', permissions: str=OUTPUT_PERMISSIONS, timeout: int=30) -> None:
|
||||
"""chmod -R <permissions> <cwd>/<path>"""
|
||||
|
||||
if not os.path.exists(os.path.join(cwd, path)):
|
||||
|
@ -538,7 +551,7 @@ def chmod_file(path, cwd='.', permissions=OUTPUT_PERMISSIONS, timeout=30):
|
|||
raise Exception('Failed to chmod {}/{}'.format(cwd, path))
|
||||
|
||||
|
||||
def chrome_args(**options):
|
||||
def chrome_args(**options) -> List[str]:
|
||||
"""helper to build up a chrome shell command with arguments"""
|
||||
|
||||
options = {**CHROME_OPTIONS, **options}
|
||||
|
|
Loading…
Reference in a new issue