util.py 1.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. import urllib.parse
  2. import urllib.request
  3. from urllib.error import URLError, HTTPError
  4. import os
  5. import re
  6. import tempfile
  7. from contextlib import contextmanager
  8. URL_REGEX = re.compile(r'http://|https://|ftp://|file://|file:\\')
  9. def is_url(filename):
  10. """Return True if string is an http or ftp path."""
  11. return isinstance(filename, str) and URL_REGEX.match(filename) is not None
  12. @contextmanager
  13. def file_or_url_context(resource_name):
  14. """Yield name of file from the given resource (i.e. file or url)."""
  15. if is_url(resource_name):
  16. url_components = urllib.parse.urlparse(resource_name)
  17. _, ext = os.path.splitext(url_components.path)
  18. try:
  19. with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as f:
  20. with urllib.request.urlopen(resource_name) as u:
  21. f.write(u.read())
  22. # f must be closed before yielding
  23. yield f.name
  24. except (URLError, HTTPError):
  25. # could not open URL
  26. os.remove(f.name)
  27. raise
  28. except (FileNotFoundError, FileExistsError, PermissionError, BaseException):
  29. # could not create temporary file
  30. raise
  31. else:
  32. os.remove(f.name)
  33. else:
  34. yield resource_name