Source code for pykern.pkio

# -*- coding: utf-8 -*-
u"""Useful I/O operations

:copyright: Copyright (c) 2015 RadiaSoft LLC.  All Rights Reserved.
from __future__ import absolute_import, division, print_function
from pykern import pkcompat
import contextlib
import copy
import errno
import glob
import io
import locale
import os
import os.path
import py
import re
import shutil
import six

#: used during unit testing see ``pykern.pkunit.save_chdir``
pkunit_prefix = None

[docs]def exception_is_not_found(exc): """True if exception is IOError and ENOENT Args: exc (BaseException): to check Returns: bool: True if is a file not found exception. """ return isinstance(exc, IOError) and exc.errno == errno.ENOENT or isinstance(exc, py.error.ENOENT)
[docs]def expand_user_path(path): """Calls expanduser on path If `pkunit_prefix` is set, will prefix, too. Args: path (str): path to expand Returns: py.path.Local: expanded path """ return py_path(path)
[docs]def has_file_extension(filename, to_check): """if matches any of the file extensions Args: filename (str|py.path.local): what to check to_check (str|tuple|list): is without '.' and lower Returns: bool: if any of the extensions matches """ if isinstance(to_check, six.string_types): to_check = (to_check) e = py_path(filename).ext[1:].lower() return e in to_check
[docs]def mkdir_parent(path): """Create the directories and their parents (if necessary) Args: path (str): dir to create Returns: py.path.local: path """ return py_path(path).ensure(dir=True)
[docs]def mkdir_parent_only(path): """Create the paths' parent directories. Args: path (str): children of dir to create Returns: py.path.local: parent directory of path """ return mkdir_parent(py_path(path).dirname)
[docs]def py_path(path=None): """Creates a py.path.Local object Will expanduser, if needed. If `pkunit_prefix` is set, will prefix, too. Args: path (str): path to convert (or None for current dir) Returns: py.path.Local: path """ global pkunit_prefix res = py.path.local(path, expanduser=True) if pkunit_prefix: # Allow for <test>_work and <test>_data so we don't add # prefix if there's a common parent directory. if not str(res).startswith(pkunit_prefix.dirname): res = pkunit_prefix.join(res) py.path.local(res.dirname).ensure(dir=True) return res
[docs]def read_text(filename): """Open file, read with preferred encoding text, and close. Args: filename (str or py.path.Local): File to open Returns: str: contest of `filename` """ fn = py_path(filename) with, encoding=locale.getpreferredencoding()) as f: return;
[docs]@contextlib.contextmanager def save_chdir(dirname, mkdir=False, is_pkunit_prefix=False): """Save current directory, change to directory, and restore. Args: dirname (str): directory to change to mkdir (bool): Make the directory? is_pkunit_prefix (bool): If True, sets pkunit_prefix. Returns: str: current directory before `chdir` """ global pkunit_prefix prev_d = py.path.local().realpath() prev_ppp = pkunit_prefix try: if is_pkunit_prefix: d = py.path.local(dirname) else: d = py_path(dirname) if mkdir and not d.check(dir=1): mkdir_parent(d) os.chdir(str(d)) if is_pkunit_prefix: pkunit_prefix = py.path.local(d) yield d.realpath() finally: os.chdir(str(prev_d)) if is_pkunit_prefix: pkunit_prefix = prev_ppp
[docs]def sorted_glob(path): """sorted list of py.path.Local objects, non-recursive Args: path (py.path.Local or str): pattern Returns: list: py.path.Local objects """ return sorted(py_path(f) for f in glob.glob(str(path)))
[docs]def unchecked_remove(*paths): """Remove files or directories, ignoring OSError. Will not remove '/' or '.' Args: paths (str): paths to remove """ cwd = py_path() for a in paths: p = py_path(a) assert len( > 1, \ '{}: will not remove root directory'.format(p) assert cwd != p, \ '{}: will not remove current directory'.format(p) try: os.remove(str(a)) except OSError: try: shutil.rmtree(str(a), ignore_errors=True) except OSError: pass
[docs]def walk_tree(dirname, file_re=None): """Return list files (only) as py.path's, top down, sorted If you want to go bottom up, just reverse the list. Args: dirname (str): directory to walk match_re (re or str): Optionally, only return files which match file_re Yields: py.path.local: paths in sorted order """ fr = file_re if fr and not hasattr(fr, 'search'): fr = re.compile(fr) dirname = py_path(dirname).realpath() dn = str(dirname) res = [] for r, d, files in os.walk(dn, topdown=True, onerror=None, followlinks=False): for f in files: p = py_path(r).join(f) if fr and not continue res.append(p) # Not an iterator, but works as one. Don't assume always will return list return sorted(res)
[docs]def write_text(filename, contents): """Open file, write text with preferred encoding, and close. Args: filename (str or py.path.Local): File to open contents (str): New contents Returns: py.path.local: `filename` as :class:`py.path.Local` """ fn = py_path(filename) with, 'w', encoding=locale.getpreferredencoding()) as f: f.write(pkcompat.locale_str(contents)) return fn