├── .circleci └── config.yml ├── .gitignore ├── .textfile ├── LICENSE.txt ├── README.rst ├── appveyor.yml ├── bin ├── __init__.py ├── cat.py ├── compat.py ├── env.py ├── errmd5.py ├── fold.py ├── head.py ├── pargs.py ├── pwd.py ├── repeat.py └── sha256sum.py ├── dev-requirements.txt ├── helper.py ├── pytest.ini ├── setup.cfg ├── setup.py ├── tests ├── __init__.py ├── test_chdir.py ├── test_commands.py ├── test_env.py ├── test_glob.py └── test_util.py └── ush.py /.circleci/config.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | shared: &shared 3 | working_directory: ~/python-ush 4 | environment: 5 | - BASH_ENV: ~/.bash_env 6 | steps: 7 | - run: 8 | name: set bash environment 9 | command: echo 'export PATH="${HOME}/.local/bin:${PATH}"' >> ~/.bash_env 10 | 11 | - checkout 12 | 13 | - run: 14 | name: make dev-requirements.txt unique for the python version 15 | command: | 16 | echo "# $(python --version)" >> dev-requirements.txt 17 | cat dev-requirements.txt 18 | 19 | - restore_cache: 20 | keys: 21 | - dependencies-{{ checksum "dev-requirements.txt" }} 22 | - dependencies- 23 | 24 | - run: 25 | name: install dependencies 26 | command: | 27 | pip install --user -r dev-requirements.txt 28 | 29 | - save_cache: 30 | paths: [ ~/.local ] 31 | key: dependencies-{{ checksum "dev-requirements.txt" }} 32 | 33 | - run: 34 | name: run tests 35 | command: | 36 | LANG=C.UTF-8 pytest --doctest-glob=*.rst -vv 37 | 38 | jobs: 39 | "Python 2.7": 40 | <<: *shared 41 | docker: [image: "circleci/python:2.7"] 42 | "Python 3.4": 43 | <<: *shared 44 | docker: [image: "circleci/python:3.4"] 45 | "Python 3.5": 46 | <<: *shared 47 | docker: [image: "circleci/python:3.5"] 48 | "Python 3.6": 49 | <<: *shared 50 | docker: [image: "circleci/python:3.6"] 51 | 52 | workflows: 53 | version: 2 54 | test: 55 | jobs: ["Python 2.7", "Python 3.4", "Python 3.5", "Python 3.6"] 56 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.egg-info/ 2 | .cache/ 3 | .stdout 4 | .stderr 5 | build/ 6 | dist/ 7 | -------------------------------------------------------------------------------- /.textfile: -------------------------------------------------------------------------------- 1 | 123 2 | 1234 3 | 12345 4 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Copyright (c) 2017 Thiago de Arruda 2 | 3 | Permission is hereby granted, free of charge, to any person 4 | obtaining a copy of this software and associated documentation 5 | files (the "Software"), to deal in the Software without 6 | restriction, including without limitation the rights to use, 7 | copy, modify, merge, publish, distribute, sublicense, and/or sell 8 | copies of the Software, and to permit persons to whom the 9 | Software is furnished to do so, subject to the following 10 | conditions: 11 | 12 | The above copyright notice and this permission notice shall be 13 | included in all copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 16 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES 17 | OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 18 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT 19 | HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, 20 | WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 21 | FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR 22 | OTHER DEALINGS IN THE SOFTWARE. 23 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | .. vim: ft=doctest 2 | ush - Unix Shell 3 | ================ 4 | 5 | .. image:: https://circleci.com/gh/tarruda/python-ush.svg?style=svg 6 | :target: https://circleci.com/gh/tarruda/python-ush 7 | 8 | .. image:: https://ci.appveyor.com/api/projects/status/p5n9fy83nx4ac24b?svg=true 9 | :target: https://ci.appveyor.com/project/tarruda/python-ush 10 | 11 | Python library that provides a convenient but powerful API for invoking external 12 | commands. Features: 13 | 14 | - Idiomatic API for invoking commands 15 | - Command chaining with `|` 16 | - Redirection 17 | - Windows/Unix support 18 | - Python2/3 support 19 | - Filename argument expansion (globbing) 20 | 21 | Installation 22 | ------------ 23 | 24 | .. code-block:: 25 | 26 | pip install ush 27 | 28 | 29 | Basic Usage 30 | ----------- 31 | 32 | >>> import os; os.environ['LANG'] = 'C.UTF-8' 33 | >>> from ush import Shell 34 | >>> sh = Shell() 35 | >>> ls = sh('ls') 36 | 37 | The ``ls`` variable is a ``Command`` object that wraps the ``ls`` external command: 38 | 39 | >>> ls 40 | ls 41 | 42 | Calling the command without arguments will invoke the command and return the 43 | exit code: 44 | 45 | >>> ls() 46 | (0,) 47 | 48 | This is how arguments can be added: 49 | 50 | >>> ls('-l', '-a', '-h') 51 | ls -l -a -h 52 | >>> ls('-lah')() 53 | (0,) 54 | 55 | Adding arguments actually creates new ``Command`` instances with the appended 56 | arguments. If the same arguments are to be used in future invocations, it can be 57 | useful to save to a variable: 58 | 59 | >>> ls = ls('--hide=__pycache__', '--hide=*.py*', '--hide=*.yml', '--hide=*.txt', env={'LANG': 'C.UTF-8'}) 60 | 61 | By default, standard input, output and error are inherited from the python 62 | process. To capture output, simply call `str()` or `unicode()` on the ``Command`` 63 | object: 64 | 65 | >>> str(ls) 66 | 'README.rst\nbin\npytest.ini\nsetup.cfg\ntests\n' 67 | 68 | ``Command`` instances are also iterable, which is useful to process commands that 69 | output a lot of data without consuming everything in memory. By default, the 70 | iterator treats the command output as utf-8 and yields one item per line: 71 | 72 | >>> files = [] 73 | >>> for line in ls: 74 | ... files.append(line) 75 | ... 76 | >>> files 77 | [u'README.rst', u'bin', u'pytest.ini', u'setup.cfg', u'tests'] 78 | 79 | It is possible to iterate on raw chunks of data (as received from the command) 80 | by calling the `iter_raw()` method. 81 | 82 | >>> list(ls.iter_raw()) 83 | [b'README.rst\nbin\npytest.ini\nsetup.cfg\ntests\n'] 84 | 85 | The normal behavior of invoking commands is return the exit code, even if it is 86 | an error: 87 | 88 | >>> ls('invalid-file')() 89 | (2,) 90 | 91 | If the command is passed `raise_on_error=True`, it will raise an exception when 92 | the external command returns non-zero codes: 93 | 94 | >>> ls('invalid-file', raise_on_error=True)() 95 | Traceback (most recent call last): 96 | ... 97 | ProcessError: One more commands failed 98 | 99 | The directory and environment of the command can be customized with the ``cwd`` 100 | and ``env`` options, respectively: 101 | 102 | >>> ls(cwd='bin', env={'LS_COLORS': 'ExGxFxdxCxDxDxhbadExEx'})() 103 | (0,) 104 | 105 | Default options 106 | --------------- 107 | 108 | ``Shell`` instances act like a factory for ``Command`` objects, and can be used to 109 | hold default options for commands created by it: 110 | 111 | >>> sh = Shell(raise_on_error=True) 112 | >>> sort, cat, echo = sh(['sort', '--reverse'], 'cat', 'echo') 113 | >>> sort 114 | sort --reverse (raise_on_error=True) 115 | 116 | It is possible to override when calling the ``Shell`` object: 117 | 118 | >>> sort = sh(['sort', '--reverse'], cwd='bin', raise_on_error=None) 119 | >>> sort 120 | sort --reverse (cwd=bin) 121 | 122 | >>> sort = sort(cwd=None) 123 | >>> sort 124 | sort --reverse 125 | 126 | Pipelines 127 | --------- 128 | 129 | Like with unix shells, it is possible to chain commands via the pipe (`|`) 130 | operator: 131 | 132 | >>> ls | sort 133 | ls --hide=__pycache__ --hide=*.py* --hide=*.yml --hide=*.txt (env={'LANG': 'C.UTF-8'}) | sort --reverse 134 | 135 | Everything that can be done with single commands, can also be done with 136 | pipelines: 137 | 138 | >>> (ls | sort)() 139 | (0, 0) 140 | >>> str(ls | sort) 141 | 'tests\nsetup.cfg\npytest.ini\nbin\nREADME.rst\n' 142 | >>> list(ls | sort) 143 | [u'tests', u'setup.cfg', u'pytest.ini', u'bin', u'README.rst'] 144 | 145 | Redirection 146 | ----------- 147 | 148 | Redirecting stdin/stdout to files is also done with the `|` operator, but 149 | chained with filenames instead of other ``Command`` instances: 150 | 151 | >>> (ls | sort | '.stdout')() 152 | (0, 0) 153 | >>> str(cat('.stdout')) 154 | 'tests\nsetup.cfg\npytest.ini\nbin\nREADME.rst\n' 155 | >>> str('setup.cfg' | cat) 156 | '[metadata]\ndescription-file = README.rst\n\n[bdist_wheel]\nuniversal=1\n' 157 | 158 | In other words, a filename on the left side of the `|` will connect the file to 159 | the command's stdin, a filename on the right side of the `|` will write the 160 | command's stdout to the file. 161 | 162 | When redirecting stdout, the file is truncated by default. To append to the 163 | file, add the `+` suffix to the filename, For example: 164 | 165 | >>> (echo('some more data') | cat | '.stdout+')() 166 | (0, 0) 167 | >>> str(cat('.stdout')) 168 | 'tests\nsetup.cfg\npytest.ini\nbin\nREADME.rst\nsome more data\n' 169 | 170 | While only the first and last command of a pipeline may redirect stdin/stdout, 171 | any command in a pipeline may redirect stderr through the ``stderr`` option: 172 | 173 | >>> ls('invalid-file', stderr='.stderr', raise_on_error=False)() 174 | (2,) 175 | >>> str(cat('.stderr')) #doctest: +SKIP 176 | "ls: cannot access 'invalid-file': No such file or directory\n" 177 | 178 | Besides redirecting to/from filenames, it is possible to redirect to/from any 179 | file-like object: 180 | 181 | >>> from six import BytesIO 182 | >>> sink = BytesIO() 183 | >>> ls('invalid-file', stderr=sink, raise_on_error=False)() 184 | (2,) 185 | >>> sink.getvalue() #doctest: +SKIP 186 | b"ls: cannot access 'invalid-file': No such file or directory\n" 187 | >>> sink = BytesIO() 188 | >>> (BytesIO(b'some in-memory data') | cat | sink)() 189 | (0,) 190 | >>> sink.getvalue() 191 | b'some in-memory data' 192 | 193 | To simplify passing strings to stdin of commands, the ``sh.echo`` helper is 194 | provided: 195 | 196 | >>> sink = BytesIO() 197 | >>> (sh.echo('some in-memory data') | cat | sink)() 198 | (0,) 199 | >>> sink.getvalue() 200 | b'some in-memory data' 201 | 202 | >>> sink = BytesIO() 203 | >>> (sh.echo(b'some in-memory data') | cat | sink)() 204 | (0,) 205 | >>> sink.getvalue() 206 | b'some in-memory data' 207 | 208 | ``sh.echo`` is just a small wrapper around ``BytesIO`` or ``StringIO``. 209 | 210 | Environment 211 | ----------- 212 | 213 | Like with `subprocess.Popen`, environment variables are inherited by default, 214 | but there are some differences with how the ``env`` option is handled: 215 | 216 | 1- The contents of the ``env`` option is merged with the current process's 217 | environment by default: 218 | 219 | >>> import os; os.environ['USH_TEST_VAR1'] = 'v1' 220 | >>> env, grep = sh('env', 'grep', env={'USH_TEST_VAR2': 'v2'}) 221 | >>> list(sorted(env(env={'USH_TEST_VAR3': 'v3'}) | grep('^USH_TEST_'))) 222 | [u'USH_TEST_VAR1=v1', u'USH_TEST_VAR2=v2', u'USH_TEST_VAR3=v3'] 223 | 224 | 2- To disable merging with the current process's environment (and adopt 225 | `subprocess.Popen` behavior), pass `merge_env=False` with the ``env`` option. 226 | 227 | >>> list(sorted(env(env={'USH_TEST_VAR3': 'v3'}, merge_env=False) | grep('^USH_TEST_'))) 228 | [u'USH_TEST_VAR2=v2', u'USH_TEST_VAR3=v3'] 229 | 230 | 3- Variables can be cleared in the child process by passing a ``None`` value. 231 | 232 | >>> list(sorted(env(env={'USH_TEST_VAR1': None}) | grep('^USH_TEST_'))) 233 | [u'USH_TEST_VAR2=v2'] 234 | 235 | As shown in the above examples, setting the ``env`` option always merges the 236 | variables with previous invocations. To clear the value of the option, simply 237 | pass ``None`` as the ``env`` option: 238 | 239 | >>> env = env(env=None) 240 | >>> list(sorted(env | grep('^USH_TEST_'))) 241 | [u'USH_TEST_VAR1=v1'] 242 | >>> env = env(env={'USH_TEST_VAR2': '2'}) 243 | >>> list(sorted(env | grep('^USH_TEST_'))) 244 | [u'USH_TEST_VAR1=v1', u'USH_TEST_VAR2=2'] 245 | 246 | 247 | Globbing 248 | -------- 249 | 250 | Arguments passed to ``Command`` instances can be subject to filename 251 | expansion. This feature is enabled with the ``glob`` option: 252 | 253 | >>> echo = echo(glob=True) 254 | >>> list(sorted(str(echo('*.py')).split())) 255 | ['helper.py', 'setup.py', 'ush.py'] 256 | 257 | To prevent messing with command switches, arguments starting with "-" are not 258 | expanded: 259 | 260 | >>> list(sorted(str(echo('-*.py')).split())) 261 | ['-*.py'] 262 | 263 | With Python 3.5+, this expansion can be recursive: 264 | 265 | >>> list(sorted(str(echo('**/__init__.py')).split())) #doctest: +SKIP 266 | ['bin/__init__.py', 'tests/__init__.py'] 267 | 268 | Expansion is done relative to the command's ``cwd``: 269 | 270 | >>> list(sorted(str(echo('**/__init__.py', cwd='bin')).split())) #doctest: +SKIP 271 | ['__init__.py'] 272 | >>> list(sorted(str(echo('../**/__init__.py', cwd='bin')).split())) #doctest: +SKIP 273 | ['../tests/__init__.py', '__init__.py'] 274 | 275 | 276 | Module syntax 277 | ------------- 278 | 279 | It is possible to export ``Shell`` instances as modules, which enables a 280 | convenient syntax for importing commands into the current namespace: 281 | 282 | >>> sh.export_as_module('mysh') 283 | >>> from ush.mysh import cat 284 | >>> str('setup.cfg' | cat) 285 | '[metadata]\ndescription-file = README.rst\n\n[bdist_wheel]\nuniversal=1\n' 286 | 287 | By default, the module name passed to ``Shell.export_as_module`` is prefixed by 288 | ``ush.``. It is possible to specify the full module name like this: 289 | 290 | >>> sh.export_as_module('mysh', full_name=True) 291 | >>> from mysh import cat 292 | 293 | Since only valid python identifiers can be imported with the module syntax, some 294 | additional work is required to import commands which are not valid identifiers. 295 | For example: 296 | 297 | >>> sh.alias(apt_get='apt-get') 298 | >>> from mysh import apt_get 299 | >>> apt_get 300 | apt-get (raise_on_error=True) 301 | 302 | A builtin ``Shell`` instance with common options and aliases is already 303 | available as the ``ush.sh`` module: 304 | 305 | >>> import ush.sh as s 306 | >>> s #doctest: +ELLIPSIS 307 | 308 | 309 | This feature is inspired by `sh.py http://amoffat.github.io/sh/`. 310 | -------------------------------------------------------------------------------- /appveyor.yml: -------------------------------------------------------------------------------- 1 | build: off 2 | deploy: off 3 | clone_depth: 1 4 | 5 | install: 6 | - "python --version" 7 | - "python -c \"import struct; print(struct.calcsize('P') * 8)\"" 8 | - "pip install --disable-pip-version-check --user --upgrade pip" 9 | - "pip install -r dev-requirements.txt" 10 | 11 | test_script: 12 | - "pytest --doctest-glob=*.txt -vv" 13 | -------------------------------------------------------------------------------- /bin/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tarruda/python-ush/7a600037a0db342ea25ff5c22d486420b8de6286/bin/__init__.py -------------------------------------------------------------------------------- /bin/cat.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import shutil 3 | 4 | from compat import stdin, stdout 5 | 6 | 7 | def parse_argv(): 8 | parser = argparse.ArgumentParser('concatenate files to stdout') 9 | parser.add_argument('files', type=argparse.FileType('rb'), nargs='*', 10 | default=[stdin]) 11 | return parser.parse_args() 12 | 13 | 14 | def main(): 15 | args = parse_argv() 16 | for f in args.files: 17 | shutil.copyfileobj(f, stdout) 18 | f.close() 19 | 20 | 21 | if __name__ == '__main__': 22 | main() 23 | -------------------------------------------------------------------------------- /bin/compat.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | if sys.version_info[0] >= 3: 4 | stdin = sys.stdin.buffer 5 | stdout = sys.stdout.buffer 6 | stderr = sys.stderr.buffer 7 | else: 8 | stdin = sys.stdin 9 | stdout = sys.stdout 10 | stderr = sys.stderr 11 | -------------------------------------------------------------------------------- /bin/env.py: -------------------------------------------------------------------------------- 1 | # Helper to test environment variable inheritance. Prints only environment 2 | # variables starting with USH_ 3 | import os 4 | 5 | from compat import stdin, stdout 6 | 7 | 8 | def main(): 9 | for k in sorted(os.environ): 10 | if k.startswith('USH_'): 11 | print('{}={}'.format(k, os.environ[k])) 12 | 13 | 14 | if __name__ == '__main__': 15 | main() 16 | -------------------------------------------------------------------------------- /bin/errmd5.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | import shutil 3 | 4 | import six 5 | 6 | from compat import stdin, stdout, stderr 7 | 8 | 9 | def main(): 10 | out = six.BytesIO() 11 | shutil.copyfileobj(stdin, out) 12 | md5 = hashlib.md5() 13 | md5.update(out.getvalue()) 14 | stdout.write(out.getvalue()) 15 | stdout.flush() 16 | stderr.write(md5.hexdigest().encode() + b'\n') 17 | stderr.flush() 18 | 19 | 20 | if __name__ == '__main__': 21 | main() 22 | -------------------------------------------------------------------------------- /bin/fold.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | 4 | from compat import stdin, stdout 5 | 6 | 7 | def parse_argv(): 8 | parser = argparse.ArgumentParser('wrap input lines of width of WIDTH') 9 | parser.add_argument('-w', '--width', type=int, required=True) 10 | return parser.parse_args() 11 | 12 | 13 | def flush(buf): 14 | stdout.write(b''.join(buf + [b'\n'])) 15 | 16 | 17 | def main(): 18 | args = parse_argv() 19 | width = args.width 20 | buf = [] 21 | while True: 22 | c = stdin.read(1) 23 | if not c: 24 | break 25 | if c == b'\n': 26 | flush(buf) 27 | buf = [] 28 | continue 29 | if len(buf) == width: 30 | flush(buf) 31 | buf = [c] 32 | continue 33 | buf.append(c) 34 | 35 | 36 | if __name__ == '__main__': 37 | main() 38 | -------------------------------------------------------------------------------- /bin/head.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | 3 | from compat import stdin, stdout 4 | 5 | 6 | def parse_argv(): 7 | parser = argparse.ArgumentParser('output first COUNT input bytes') 8 | parser.add_argument('-c', '--count', type=int, required=True) 9 | return parser.parse_args() 10 | 11 | 12 | def main(): 13 | args = parse_argv() 14 | remaining = args.count 15 | while remaining: 16 | chunk = stdin.read(1024) 17 | if not chunk: 18 | break 19 | output = chunk[:remaining] if remaining < len(chunk) else chunk 20 | remaining -= len(output) 21 | stdout.write(output) 22 | 23 | 24 | if __name__ == '__main__': 25 | main() 26 | -------------------------------------------------------------------------------- /bin/pargs.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | 4 | def main(): 5 | for arg in sys.argv[1:]: 6 | print(arg) 7 | 8 | 9 | if __name__ == '__main__': 10 | main() 11 | -------------------------------------------------------------------------------- /bin/pwd.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | def main(): 4 | print(os.path.abspath(os.curdir)) 5 | 6 | if __name__ == '__main__': 7 | main() 8 | -------------------------------------------------------------------------------- /bin/repeat.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import sys 3 | 4 | from compat import stdout 5 | 6 | 7 | if sys.version_info[0] < 3: 8 | range = xrange 9 | 10 | 11 | def parse_argv(): 12 | parser = argparse.ArgumentParser('repeat argument') 13 | parser.add_argument('data') 14 | parser.add_argument('-c', '--count', type=int, required=True) 15 | return parser.parse_args() 16 | 17 | 18 | def main(): 19 | args = parse_argv() 20 | data = args.data.encode('utf8') 21 | for _ in range(args.count): 22 | stdout.write(data) 23 | 24 | 25 | if __name__ == '__main__': 26 | main() 27 | -------------------------------------------------------------------------------- /bin/sha256sum.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | 3 | 4 | from compat import stdin, stdout 5 | 6 | 7 | def main(): 8 | m = hashlib.sha256() 9 | while True: 10 | chunk = stdin.read(1024) 11 | if not chunk: 12 | break 13 | m.update(chunk) 14 | stdout.write(m.hexdigest().encode() + b'\n') 15 | 16 | 17 | if __name__ == '__main__': 18 | main() 19 | -------------------------------------------------------------------------------- /dev-requirements.txt: -------------------------------------------------------------------------------- 1 | pytest 2 | six 3 | -------------------------------------------------------------------------------- /helper.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import ush 4 | 5 | __all__ = ('cat', 'echo', 'env', 'fold', 'head', 'repeat', 'sha256sum', 6 | 'errmd5', 'pargs', 'pwd', 'STDOUT', 'PIPE', 's', 'sh') 7 | 8 | SOURCE_ROOT = os.path.join(os.path.abspath(os.path.dirname(__file__))) 9 | TEST_BIN_DIR = os.path.join(SOURCE_ROOT, 'bin') 10 | 11 | def commands(*names): 12 | argvs = [] 13 | for name in names: 14 | script = os.path.join(TEST_BIN_DIR, '{0}.py'.format(name)) 15 | argvs.append([sys.executable, script]) 16 | return sh(*argvs) 17 | 18 | 19 | def s(obj): 20 | """Helper to normalize linefeeds in strings.""" 21 | if isinstance(obj, bytes): 22 | return obj.replace(b'\n', os.linesep.encode()) 23 | else: 24 | return obj.replace('\n', os.linesep) 25 | 26 | 27 | ush.Shell().export_as_module('sh', full_name=True) 28 | import sh 29 | for name in ['cat', 'env', 'fold', 'head', 'repeat', 'sha256sum', 'errmd5', 30 | 'pargs', 'pwd']: 31 | script = os.path.join(TEST_BIN_DIR, '{0}.py'.format(name)) 32 | alias_dict = {name: [sys.executable, script]} 33 | sh.alias(**alias_dict) 34 | from sh import (cat, echo, env, fold, head, repeat, sha256sum, errmd5, pargs, 35 | pwd) 36 | 37 | STDOUT = ush.STDOUT 38 | PIPE = ush.PIPE 39 | -------------------------------------------------------------------------------- /pytest.ini: -------------------------------------------------------------------------------- 1 | [pytest] 2 | timeout = 3 3 | timeout_method = thread 4 | doctest_optionflags = NORMALIZE_WHITESPACE IGNORE_EXCEPTION_DETAIL ALLOW_UNICODE ALLOW_BYTES 5 | 6 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | description-file = README.rst 3 | 4 | [bdist_wheel] 5 | universal=1 6 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | 4 | VERSION = '4.1.0' 5 | REPO = 'https://github.com/tarruda/python-ush' 6 | 7 | setup( 8 | name='ush', 9 | py_modules=['ush'], 10 | version=VERSION, 11 | description='Powerful API for invoking with external commands', 12 | author='Thiago de Arruda', 13 | author_email='tpadilha84@gmail.com', 14 | license='MIT', 15 | url=REPO, 16 | download_url='{0}/archive/{1}.tar.gz'.format(REPO, VERSION), 17 | keywords=['sh', 'unix', 'bash', 'shell', 'glob'] 18 | ) 19 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tarruda/python-ush/7a600037a0db342ea25ff5c22d486420b8de6286/tests/__init__.py -------------------------------------------------------------------------------- /tests/test_chdir.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | import pytest 4 | 5 | from helper import pwd, s, sh 6 | 7 | 8 | @pytest.fixture() 9 | def dirs(tmpdir): 10 | gp = tmpdir.mkdir('gp') 11 | p1 = gp.mkdir('p1') 12 | c1 = p1.mkdir('c1') 13 | c2 = p1.mkdir('c2') 14 | p2 = gp.mkdir('p2') 15 | l = [gp, p1, c1, c2, p2] 16 | return [os.path.dirname(str(gp))] + [os.path.basename(str(d)) for d in l] 17 | 18 | 19 | def test_chdir(dirs): 20 | base, gp, p1, c1, c2, p2 = dirs 21 | with sh.chdir(base): 22 | assert list(pwd)[0] == base 23 | with sh.chdir(gp): 24 | assert list(pwd)[0] == os.path.join(base, gp) 25 | with sh.chdir(p1): 26 | assert list(pwd)[0] == os.path.join(base, gp, p1) 27 | with sh.chdir(c1): 28 | assert list(pwd)[0] == os.path.join(base, gp, p1, c1) 29 | with sh.chdir(c2): 30 | assert list(pwd)[0] == os.path.join(base, gp, p1, c2) 31 | assert list(pwd)[0] == os.path.join(base, gp, p1) 32 | with sh.chdir(p2): 33 | assert list(pwd)[0] == os.path.join(base, gp, p2) 34 | assert list(pwd)[0] == base 35 | -------------------------------------------------------------------------------- /tests/test_commands.py: -------------------------------------------------------------------------------- 1 | import hashlib 2 | import os 3 | 4 | import pytest 5 | from six import BytesIO, StringIO, PY2 6 | 7 | from helper import * 8 | 9 | repeat_hex = repeat('-c', '100', '0123456789abcdef') 10 | 11 | 12 | def test_simple_success(): 13 | assert cat('.textfile')() == (0,) 14 | 15 | 16 | def test_simple_failure(): 17 | assert cat('inexistent-file')() != (0,) 18 | 19 | 20 | def test_redirect_stdout(): 21 | sink = BytesIO() 22 | assert (cat('.textfile') | sink)() == (0,) 23 | assert sink.getvalue() == s(b'123\n1234\n12345\n') 24 | 25 | 26 | def test_redirect_stdin(): 27 | source = BytesIO() 28 | source.write(s(b'abc\ndef')) 29 | assert (source | cat)() == (0,) 30 | 31 | 32 | def test_redirect_stdout_and_stdin(): 33 | source = BytesIO() 34 | source.write(s(b'abc\ndef')) 35 | source.seek(0) 36 | sink = BytesIO() 37 | assert (source | cat | sink)() == (0,) 38 | assert sink.getvalue() == source.getvalue() 39 | 40 | 41 | def test_one_pipe(): 42 | sink = BytesIO() 43 | assert (repeat_hex | sha256sum | sink)() == (0, 0,) 44 | assert ( 45 | sink.getvalue() == 46 | s(b'1f1a5c83e53c9faa87badd5d17c45ffec' 47 | b'49b137430c9817dd5c9420fd96aaa3e\n')) 48 | 49 | 50 | def test_two_pipes(): 51 | sink = BytesIO() 52 | assert (repeat_hex | sha256sum | fold('-w', 16) | sink)() == (0, 0, 0,) 53 | assert ( 54 | sink.getvalue() == 55 | s(b'1f1a5c83e53c9faa\n' 56 | b'87badd5d17c45ffe\n' 57 | b'c49b137430c9817d\n' 58 | b'd5c9420fd96aaa3e\n')) 59 | 60 | 61 | def test_three_pipes(): 62 | sink = BytesIO() 63 | assert (repeat_hex | sha256sum | fold('-w', 16) | head('-c', 18) | 64 | sink)() == (0, 0, 0, 0,) 65 | assert sink.getvalue() == s(b'1f1a5c83e53c9faa\n8') 66 | 67 | 68 | def test_stderr_redirect(): 69 | stderr_sink = BytesIO() 70 | sink = BytesIO() 71 | source = BytesIO() 72 | source.write(b'123\n') 73 | source.seek(0) 74 | assert (source | errmd5(stderr=stderr_sink) | errmd5(stderr=stderr_sink) | 75 | sink)() == (0, 0) 76 | assert stderr_sink.getvalue() == ( 77 | s(b'ba1f2511fc30423bdbb183fe33f3dd0f\n') * 2) 78 | assert sink.getvalue() == s(b'123\n') 79 | 80 | 81 | def test_stderr_redirect_to_stdout(): 82 | sink = BytesIO() 83 | source = BytesIO() 84 | source.write(s(b'123\n')) 85 | source.seek(0) 86 | assert (source | errmd5(stderr=STDOUT) | errmd5 | sink)() == (0, 0) 87 | assert sink.getvalue() == s(b'123\n' 88 | b'ba1f2511fc30423bdbb183fe33f3dd0f\n') 89 | source.seek(0) 90 | sink.seek(0) 91 | assert (source | errmd5(stderr=STDOUT) | errmd5(stderr=STDOUT) | 92 | sink)() == (0, 0) 93 | assert sink.getvalue() == s(b'123\n' 94 | b'ba1f2511fc30423bdbb183fe33f3dd0f\n' 95 | b'ba5d6480bba42f55a708ac7096374f7a\n') 96 | 97 | 98 | def test_string_input_output(): 99 | assert str(repeat('-c', '5', '123')) == '123123123123123' 100 | assert bytes(repeat('-c', '5', '123')) == b'123123123123123' 101 | if PY2: 102 | assert unicode(repeat('-c', '5', '123')) == u'123123123123123' 103 | assert str(echo(s('abc\ndef')) | cat) == s('abc\ndef') 104 | assert bytes(echo(s('abc\ndef')) | cat) == s(b'abc\ndef') 105 | if PY2: 106 | assert unicode(echo(s('abc\ndef')) | cat) == s(u'abc\ndef') 107 | 108 | 109 | def test_stdin_redirect_file(): 110 | assert str('.textfile' | cat) == s('123\n1234\n12345\n') 111 | 112 | 113 | def test_stdout_stderr_redirect_file(): 114 | (echo('hello') | errmd5(stderr='.stderr') | '.stdout')() 115 | with open('.stdout', 'rb') as f: 116 | assert f.read() == b'hello' 117 | with open('.stderr', 'rb') as f: 118 | assert f.read() == s(b'5d41402abc4b2a76b9719d911017c592\n') 119 | (echo(s('\nworld\n')) | errmd5(stderr='.stderr+') | '.stdout+')() 120 | with open('.stdout', 'rb') as f: 121 | assert f.read() == s(b'hello\nworld\n') 122 | with open('.stderr', 'rb') as f: 123 | assert f.read() == s(b'5d41402abc4b2a76b9719d911017c592\n' 124 | b'81f82f69f5be2752005dae73e0f22f76\n') 125 | 126 | 127 | def test_iterator_input(): 128 | assert str((n for n in [1, 2, 3, 4]) | cat) == '1234' 129 | assert str(['ab', 2, s('\n'), 5] | cat) == s('ab2\n5') 130 | 131 | 132 | def test_iterator_output(): 133 | chunks = [] 134 | for chunk in repeat('-c', '5', '123'): 135 | chunks.append(chunk) 136 | assert chunks == ['123123123123123'] 137 | chunks = [] 138 | for chunk in (echo(s(b'123\n')) | errmd5(stderr=STDOUT) | 139 | errmd5(stderr=STDOUT)): 140 | chunks.append(chunk) 141 | assert chunks == [ 142 | '123', 143 | 'ba1f2511fc30423bdbb183fe33f3dd0f', 144 | 'ba5d6480bba42f55a708ac7096374f7a', 145 | ] 146 | 147 | 148 | def test_iterator_output_multiple_pipes(): 149 | chunks = [] 150 | for chunk in (echo(s(b'123\n')) | errmd5(stderr=PIPE) | 151 | errmd5(stderr=PIPE)): 152 | chunks.append(chunk) 153 | assert len(chunks) == 3 154 | assert (s('ba1f2511fc30423bdbb183fe33f3dd0f'), None, None) in chunks 155 | assert (None, s('ba1f2511fc30423bdbb183fe33f3dd0f'), None) in chunks 156 | assert (None, None, s('123')) in chunks 157 | 158 | 159 | @pytest.mark.parametrize('chunk_factor', [16, 32, 64, 128, 256]) 160 | def test_big_data(chunk_factor): 161 | def generator(): 162 | b = ( 163 | b'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/' 164 | * chunk_factor 165 | ) 166 | sent = 0 167 | total = 32 * 1024 * 1024 168 | while sent < total: 169 | yield b 170 | sent += len(b) 171 | md5 = hashlib.md5() 172 | stderr_1 = None 173 | stderr_2 = None 174 | chunk_count = 0 175 | for err1, err2, chunk in ( 176 | generator() | errmd5(stderr=PIPE) | errmd5(stderr=PIPE)).iter_raw(): 177 | chunk_count += 1 178 | if err1 is not None: 179 | assert stderr_1 is None 180 | stderr_1 = err1 181 | elif err2 is not None: 182 | assert stderr_2 is None 183 | stderr_2 = err2 184 | else: 185 | md5.update(chunk) 186 | assert stderr_1 == s(b'80365aea26be3a31ce7f953d7b01ea0d\n') 187 | assert stderr_2 == s(b'80365aea26be3a31ce7f953d7b01ea0d\n') 188 | assert md5.hexdigest() == '80365aea26be3a31ce7f953d7b01ea0d' 189 | -------------------------------------------------------------------------------- /tests/test_env.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from helper import env, s, sh 4 | 5 | 6 | def setup_function(function): 7 | os.environ.clear() 8 | os.environ['USH_VAR1'] = 'var1' 9 | os.environ['USH_VAR2'] = 'var2' 10 | 11 | 12 | def test_env_inherited_by_default(): 13 | assert str(env) == s('USH_VAR1=var1\nUSH_VAR2=var2\n') 14 | 15 | 16 | def test_new_env_merged_by_default(): 17 | assert str(env(env={'USH_VAR3': 'var3', 'USH_VAR4': 'var4'})) == s( 18 | 'USH_VAR1=var1\nUSH_VAR2=var2\nUSH_VAR3=var3\nUSH_VAR4=var4\n') 19 | 20 | 21 | def test_env_override(): 22 | assert str(env(merge_env=False)) == s('USH_VAR1=var1\nUSH_VAR2=var2\n') 23 | assert str(env(merge_env=False, env={})) == '' 24 | assert str(env(merge_env=False, env={ 25 | 'USH_VAR3': 'var3', 'USH_VAR4': 'var4'})) == s( 26 | 'USH_VAR3=var3\nUSH_VAR4=var4\n') 27 | 28 | 29 | def test_unset_env(): 30 | assert str(env(env={'USH_VAR1': None})) == s('USH_VAR2=var2\n') 31 | assert str(env(env={'USH_VAR1': None, 'USH_VAR2': None})) == '' 32 | 33 | 34 | def test_shell_setenv(): 35 | with sh.setenv({'USH_VAR9': 'var9'}): 36 | assert str(env) == s('USH_VAR1=var1\nUSH_VAR2=var2\nUSH_VAR9=var9\n') 37 | with sh.setenv({'USH_VAR2': None}): 38 | assert str(env) == s('USH_VAR1=var1\nUSH_VAR9=var9\n') 39 | with sh.setenv({'USH_VAR9': None}): 40 | assert str(env) == s('USH_VAR1=var1\n') 41 | assert str(env) == s('USH_VAR1=var1\nUSH_VAR9=var9\n') 42 | assert str(env) == s('USH_VAR1=var1\nUSH_VAR2=var2\nUSH_VAR9=var9\n') 43 | assert str(env) == s('USH_VAR1=var1\nUSH_VAR2=var2\n') 44 | 45 | -------------------------------------------------------------------------------- /tests/test_glob.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | from helper import pargs, s 5 | 6 | 7 | def norm_seps(paths): 8 | return [ 9 | p.replace('/', os.path.sep) for p in paths 10 | ] 11 | 12 | 13 | def test_glob_disabled_by_default(): 14 | assert list(sorted(pargs('*.py'))) == ['*.py'] 15 | 16 | 17 | def test_glob(): 18 | assert list(sorted( 19 | pargs('*.py', glob=True))) == [ 20 | 'helper.py', 'setup.py', 'ush.py' 21 | ] 22 | 23 | 24 | def test_dont_expand_leading_dash(): 25 | assert list(sorted(pargs('-*.py', glob=True))) == ['-*.py'] 26 | 27 | 28 | def test_glob_with_relative_dir(): 29 | assert list(sorted( 30 | pargs('../*.py', cwd='bin', glob=True))) == norm_seps([ 31 | '../helper.py', '../setup.py', '../ush.py' 32 | ]) 33 | 34 | 35 | if sys.version_info >= (3, 5): 36 | def test_glob_recursive(): 37 | assert list(sorted( 38 | pargs('**/*.py', glob=True))) == norm_seps([ 39 | 'bin/__init__.py', 40 | 'bin/cat.py', 41 | 'bin/compat.py', 42 | 'bin/env.py', 43 | 'bin/errmd5.py', 44 | 'bin/fold.py', 45 | 'bin/head.py', 46 | 'bin/pargs.py', 47 | 'bin/pwd.py', 48 | 'bin/repeat.py', 49 | 'bin/sha256sum.py', 50 | 'helper.py', 51 | 'setup.py', 52 | 'tests/__init__.py', 53 | 'tests/test_chdir.py', 54 | 'tests/test_commands.py', 55 | 'tests/test_env.py', 56 | 'tests/test_glob.py', 57 | 'tests/test_util.py', 58 | 'ush.py' 59 | ]) 60 | 61 | 62 | def test_glob_recursive_with_relative_dir(): 63 | assert list(sorted( 64 | pargs('../**/*.py', cwd='bin', glob=True))) == norm_seps([ 65 | '../helper.py', 66 | '../setup.py', 67 | '../tests/__init__.py', 68 | '../tests/test_chdir.py', 69 | '../tests/test_commands.py', 70 | '../tests/test_env.py', 71 | '../tests/test_glob.py', 72 | '../tests/test_util.py', 73 | '../ush.py', 74 | '__init__.py', 75 | 'cat.py', 76 | 'compat.py', 77 | 'env.py', 78 | 'errmd5.py', 79 | 'fold.py', 80 | 'head.py', 81 | 'pargs.py', 82 | 'pwd.py', 83 | 'repeat.py', 84 | 'sha256sum.py' 85 | ]) 86 | -------------------------------------------------------------------------------- /tests/test_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pytest 3 | 4 | from ush import iterate_lines 5 | 6 | 7 | def chunk_iterator(data, chunk_size): 8 | while data: 9 | yield data[:chunk_size], 0 10 | data = data[chunk_size:] 11 | 12 | 13 | DATA = 'Lorem {0}ipsum dolor {0}sit amet{0}'.format(os.linesep).encode('utf-8') 14 | 15 | 16 | @pytest.mark.parametrize('chunk_size', list(range(1, len(DATA)))) 17 | def test_iterate_lines(chunk_size): 18 | lines = [l for l, i in iterate_lines(chunk_iterator(DATA, chunk_size))] 19 | assert lines == ['Lorem ', 'ipsum dolor ', 'sit amet', ''] 20 | 21 | -------------------------------------------------------------------------------- /ush.py: -------------------------------------------------------------------------------- 1 | import collections 2 | import contextlib 3 | import errno 4 | import glob 5 | import os 6 | import re 7 | import subprocess 8 | import sys 9 | import types 10 | 11 | 12 | __all__ = ('Shell', 'Command', 'InvalidPipeline', 'AlreadyRedirected', 13 | 'ProcessError') 14 | 15 | 16 | STDOUT = subprocess.STDOUT 17 | PIPE = subprocess.PIPE 18 | # Some opts have None as a valid value, so we use EMPTY as a default value when 19 | # reading opts to determine if an option was passed. 20 | EMPTY = object() 21 | # Cross/platform /dev/null specifier alias 22 | NULL = os.devnull 23 | MAX_CHUNK_SIZE = 0xffff 24 | GLOB_PATTERNS = re.compile(r'(?:\*|\?|\[[^\]]+\])') 25 | GLOB_OPTS = {} 26 | 27 | # We have python2/3 compatibility, but don't want to rely on `six` package so 28 | # this script can be used independently. 29 | try: 30 | xrange 31 | from Queue import Queue, Empty 32 | import StringIO 33 | StringIO = BytesIO = StringIO.StringIO 34 | def is_string(o): 35 | return isinstance(o, basestring) 36 | to_cstr = str 37 | PY3 = False 38 | except NameError: 39 | xrange = range 40 | from queue import Queue, Empty 41 | import io 42 | StringIO = io.StringIO 43 | BytesIO = io.BytesIO 44 | def is_string(o): 45 | return isinstance(o, str) or isinstance(o, bytes) 46 | def to_cstr(obj): 47 | if isinstance(obj, bytes): 48 | return obj 49 | return str(obj).encode('utf-8') 50 | PY3 = True 51 | if sys.version_info >= (3, 5): 52 | GLOB_OPTS = {'recursive': True} 53 | 54 | 55 | if sys.platform == 'win32': 56 | import threading 57 | def set_extra_popen_opts(opts): 58 | pass 59 | def concurrent_communicate(proc, read_streams): 60 | return concurrent_communicate_with_threads(proc, read_streams) 61 | else: 62 | import select 63 | from signal import signal, SIGPIPE, SIG_DFL 64 | _PIPE_BUF = getattr(select, 'PIPE_BUF', 512) 65 | def set_extra_popen_opts(opts): 66 | user_preexec_fn = opts.get('preexec_fn', None) 67 | def preexec_fn(): 68 | if user_preexec_fn: 69 | user_preexec_fn() 70 | # Restore SIGPIPE default handler when forked. This is required for 71 | # handling pipelines correctly. 72 | signal(SIGPIPE, SIG_DFL) 73 | opts['preexec_fn'] = preexec_fn 74 | def concurrent_communicate(proc, read_streams): 75 | return concurrent_communicate_with_select(proc, read_streams) 76 | 77 | 78 | class InvalidPipeline(Exception): 79 | pass 80 | 81 | 82 | class AlreadyRedirected(Exception): 83 | pass 84 | 85 | 86 | class ProcessError(Exception): 87 | def __init__(self, process_info): 88 | msg = 'One or more commands failed: {}'.format(process_info) 89 | super(ProcessError, self).__init__(msg) 90 | self.process_info = process_info 91 | 92 | 93 | def expand_filenames(argv, cwd): 94 | def expand_arg(arg): 95 | return [os.path.relpath(p, cwd) 96 | for p in glob.iglob(os.path.join(cwd, arg), **GLOB_OPTS)] 97 | rv = [argv[0]] 98 | for arg in argv[1:]: 99 | if arg and arg[0] != '-' and GLOB_PATTERNS.search(arg): 100 | rv += expand_arg(arg) 101 | else: 102 | rv.append(arg) 103 | return rv 104 | 105 | 106 | def update_opts_env(opts, extra_env): 107 | if extra_env is None: 108 | del opts['env'] 109 | return 110 | env = opts.get('env', None) 111 | if env is None: 112 | env = {} 113 | else: 114 | env = env.copy() 115 | env.update(extra_env) 116 | opts['env'] = env 117 | 118 | 119 | def set_environment(proc_opts): 120 | env = proc_opts.get('env', None) 121 | if env is None: 122 | return 123 | new_env = {} 124 | if proc_opts.get('merge_env', True): 125 | new_env.update(os.environ) 126 | new_env.update(env) 127 | # unset environment variables set to `None` 128 | for k in list(new_env.keys()): 129 | if new_env[k] is None: del new_env[k] 130 | proc_opts['env'] = new_env 131 | 132 | 133 | def fileobj_has_fileno(fileobj): 134 | try: 135 | fileobj.fileno() 136 | return True 137 | except Exception: 138 | return False 139 | 140 | 141 | def remove_invalid_opts(opts): 142 | new_opts = {} 143 | new_opts.update(opts) 144 | for opt in ('raise_on_error', 'merge_env', 'glob'): 145 | if opt in new_opts: del new_opts[opt] 146 | return new_opts 147 | 148 | 149 | LS = os.linesep 150 | LS_LEN = len(LS) 151 | 152 | def iterate_lines(chunk_iterator, trim_trailing_lf=False): 153 | remaining = {} 154 | for chunk, stream_id in chunk_iterator: 155 | chunk = remaining.get(stream_id, '') + chunk.decode('utf-8') 156 | last_ls_index = -LS_LEN 157 | while True: 158 | start = last_ls_index + LS_LEN 159 | try: 160 | ls_index = chunk.index(LS, start) 161 | except ValueError: 162 | remaining[stream_id] = chunk[last_ls_index + LS_LEN:] 163 | break 164 | yield chunk[start:ls_index], stream_id 165 | remaining[stream_id] = chunk[ls_index + LS_LEN:] 166 | last_ls_index = ls_index 167 | for stream_id in remaining: 168 | line = remaining[stream_id] 169 | if line or not trim_trailing_lf: 170 | yield line, stream_id 171 | 172 | 173 | def validate_pipeline(commands): 174 | for index, command in enumerate(commands): 175 | is_first = index == 0 176 | is_last = index == len(commands) - 1 177 | if not is_first and command.opts.get('stdin', None) is not None: 178 | msg = ( 179 | 'Command {0} is not the first in the pipeline and has ' 180 | 'stdin set to a value different than "None"' 181 | ).format(command) 182 | raise InvalidPipeline(msg) 183 | if not is_last and command.opts.get('stdout', None) is not None: 184 | msg = ( 185 | 'Command {0} is not the last in the pipeline and has ' 186 | 'stdout set to a value different than "None"' 187 | ).format(command) 188 | raise InvalidPipeline(msg) 189 | 190 | 191 | def wait(procs, raise_on_error): 192 | status_codes = [] 193 | result = tuple(iterate_outputs(procs, raise_on_error, status_codes)) 194 | assert result == () 195 | return tuple(status_codes) 196 | 197 | 198 | def iterate_outputs(procs, raise_on_error, status_codes): 199 | read_streams = [proc.stderr_stream for proc in procs if proc.stderr] 200 | if procs[-1].stdout: 201 | read_streams.append(procs[-1].stdout_stream) 202 | write_stream = procs[0].stdin_stream if procs[0].stdin else None 203 | co = communicate(procs) 204 | wchunk = None 205 | while True: 206 | try: 207 | ri = co.send(wchunk) 208 | if ri: 209 | rchunk, i = ri 210 | if read_streams[i]: 211 | read_streams[i].write(rchunk) 212 | else: 213 | yield ri 214 | except StopIteration: 215 | break 216 | try: 217 | wchunk = next(write_stream) if write_stream else None 218 | except StopIteration: 219 | wchunk = None 220 | status_codes += [proc.wait() for proc in procs] 221 | if raise_on_error and len(list(filter(lambda c: c != 0, status_codes))): 222 | process_info = [ 223 | (proc.argv, proc.pid, proc.returncode) for proc in procs 224 | ] 225 | raise ProcessError(process_info) 226 | 227 | 228 | def write_chunk(proc, chunk): 229 | try: 230 | proc.stdin.write(to_cstr(chunk)) 231 | except IOError as e: 232 | if e.errno == errno.EPIPE: 233 | # communicate() should ignore broken pipe error 234 | pass 235 | elif (e.errno == errno.EINVAL and proc.poll() is not None): 236 | # Issue #19612: stdin.write() fails with EINVAL 237 | # if the process already exited before the write 238 | pass 239 | else: 240 | raise 241 | 242 | 243 | def communicate(procs): 244 | # make a list of (readable streams, sinks) tuples 245 | read_streams = [proc.stderr for proc in procs if proc.stderr] 246 | if procs[-1].stdout: 247 | read_streams.append(procs[-1].stdout) 248 | writer = procs[0] 249 | if len(read_streams + [w for w in [writer] if w.stdin]) > 1: 250 | return concurrent_communicate(writer, read_streams) 251 | if writer.stdin or len(read_streams) == 1: 252 | return simple_communicate(writer, read_streams) 253 | else: 254 | return stub_communicate() 255 | 256 | 257 | def stub_communicate(): 258 | return 259 | yield 260 | 261 | 262 | def simple_communicate(proc, read_streams): 263 | if proc.stdin: 264 | while True: 265 | chunk = yield 266 | if not chunk: 267 | break 268 | write_chunk(proc, chunk) 269 | proc.stdin.close() 270 | else: 271 | read_stream = read_streams[0] 272 | while True: 273 | chunk = read_stream.read(MAX_CHUNK_SIZE) 274 | if not chunk: 275 | break 276 | yield (chunk, 0) 277 | 278 | 279 | def concurrent_communicate_with_select(proc, read_streams): 280 | reading = [] + read_streams 281 | writing = [proc.stdin] if proc.stdin else [] 282 | indexes = dict((r.fileno(), i) for i, r in enumerate(read_streams)) 283 | write_queue = collections.deque() 284 | 285 | while reading or writing: 286 | try: 287 | rlist, wlist, xlist = select.select(reading, writing, []) 288 | except select.error as e: 289 | if e.args[0] == errno.EINTR: 290 | continue 291 | raise 292 | 293 | for rstream in rlist: 294 | rchunk = os.read(rstream.fileno(), MAX_CHUNK_SIZE) 295 | if not rchunk: 296 | rstream.close() 297 | reading.remove(rstream) 298 | continue 299 | write_queue.append((yield rchunk, indexes[rstream.fileno()])) 300 | 301 | if not write_queue: 302 | write_queue.append((yield)) 303 | 304 | if not wlist: 305 | continue 306 | 307 | while write_queue: 308 | wchunk = write_queue.popleft() 309 | if wchunk is None: 310 | assert not write_queue 311 | writing = [] 312 | proc.stdin.close() 313 | break 314 | wchunk = to_cstr(wchunk) 315 | chunk = wchunk[:_PIPE_BUF] 316 | if len(wchunk) > _PIPE_BUF: 317 | write_queue.appendleft(wchunk[_PIPE_BUF:]) 318 | try: 319 | written = os.write(proc.stdin.fileno(), chunk) 320 | except OSError as e: 321 | if e.errno != errno.EPIPE: 322 | raise 323 | writing = [] 324 | proc.stdin.close() 325 | else: 326 | if len(chunk) > written: 327 | write_queue.appendleft(chunk[written:]) 328 | # break so we wait for the pipe buffer to be drained 329 | break 330 | 331 | 332 | def concurrent_communicate_with_threads(proc, read_streams): 333 | def read(queue, read_stream, index): 334 | while True: 335 | chunk = read_stream.read(MAX_CHUNK_SIZE) 336 | if not chunk: 337 | break 338 | queue.put((chunk, index)) 339 | queue.put((None, index)) 340 | 341 | def write(queue, proc): 342 | while True: 343 | chunk = queue.get() 344 | if not chunk: 345 | break 346 | write_chunk(proc, chunk) 347 | proc.stdin.close() 348 | 349 | rqueue = Queue(maxsize=1) 350 | wqueue = Queue() 351 | threads = [] 352 | for i, rs in enumerate(read_streams): 353 | threads.append(threading.Thread(target=read, args=(rqueue, rs, i))) 354 | threads[-1].setDaemon(True) 355 | threads[-1].start() 356 | if proc.stdin: 357 | threads.append(threading.Thread(target=write, args=(wqueue, proc))) 358 | threads[-1].setDaemon(True) 359 | threads[-1].start() 360 | 361 | writing = True 362 | reading = len(read_streams) 363 | while writing or reading or rqueue.qsize(): 364 | if reading or rqueue.qsize(): 365 | try: 366 | rchunk, index = rqueue.get(block=not writing) 367 | if rchunk: 368 | wchunk = yield rchunk, index 369 | else: 370 | reading -= 1 371 | continue 372 | except Empty: 373 | wchunk = yield 374 | else: 375 | wchunk = yield 376 | if writing: 377 | wqueue.put(wchunk) 378 | writing = wchunk is not None 379 | for t in threads: 380 | t.join() 381 | 382 | 383 | def setup_redirect(proc_opts, key): 384 | stream = proc_opts.get(key, None) 385 | if stream in (None, STDOUT, PIPE) or fileobj_has_fileno(stream): 386 | # Simple case which will be handled automatically by Popen: stream is 387 | # STDOUT/PIPE or a file object backed by file. 388 | return None, False 389 | if is_string(stream): 390 | # stream is a string representing a filename, we'll open the file with 391 | # appropriate mode which will be set to proc_opts[key]. 392 | if key == 'stdin': 393 | proc_opts[key] = open(stream, 'rb') 394 | else: 395 | if stream.endswith('+'): 396 | proc_opts[key] = open(stream[:-1], 'ab') 397 | # On MS Windows we need to explicitly the file position to the 398 | # end or the file contents will be replaced. 399 | proc_opts[key].seek(0, os.SEEK_END) 400 | else: 401 | proc_opts[key] = open(stream, 'wb') 402 | return None, True 403 | if key == 'stdin': 404 | if hasattr(stream, 'read'): 405 | # replace with an iterator that yields data in up to 64k chunks. 406 | # This is done to avoid the yield-by-line logic when iterating 407 | # file-like objects that contain binary data. 408 | stream = fileobj_to_iterator(stream) 409 | elif hasattr(stream, '__iter__'): 410 | stream = iter(stream) 411 | proc_opts[key] = PIPE 412 | return stream, False 413 | 414 | 415 | def fileobj_to_iterator(fobj): 416 | def iterator(): 417 | while True: 418 | data = fobj.read(MAX_CHUNK_SIZE) 419 | if not data: 420 | break 421 | yield data 422 | return iterator() 423 | 424 | 425 | def echo(s): 426 | if isinstance(s, str): 427 | return StringIO(s) 428 | else: 429 | assert isinstance(s, bytes) 430 | return BytesIO(s) 431 | 432 | 433 | class RunningProcess(object): 434 | def __init__(self, popen, stdin_stream, stdout_stream, stderr_stream, 435 | argv): 436 | self.popen = popen 437 | self.stdin_stream = stdin_stream 438 | self.stdout_stream = stdout_stream 439 | self.stderr_stream = stderr_stream 440 | self.argv = argv 441 | 442 | @property 443 | def returncode(self): 444 | return self.popen.returncode 445 | 446 | @property 447 | def stdin(self): 448 | return self.popen.stdin 449 | 450 | @property 451 | def stdout(self): 452 | return self.popen.stdout 453 | 454 | @property 455 | def stderr(self): 456 | return self.popen.stderr 457 | 458 | @property 459 | def pid(self): 460 | return self.popen.pid 461 | 462 | def wait(self): 463 | return self.popen.wait() 464 | 465 | def poll(self): 466 | return self.popen.poll() 467 | 468 | 469 | class Shell(object): 470 | def __init__(self, **defaults): 471 | self.aliases = {} 472 | self.envstack = [] 473 | self.dirstack = [] 474 | if 'env' in defaults: 475 | self.envstack.append(defaults['env']) 476 | del defaults['env'] 477 | if 'cwd' in defaults: 478 | self.dirstack.append(defaults['cwd']) 479 | del defaults['cwd'] 480 | self.defaults = defaults 481 | self.echo = echo 482 | 483 | def __call__(self, *argvs, **opts): 484 | rv = [] 485 | for argv in argvs: 486 | if is_string(argv): 487 | argv = self.aliases.get(argv, argv) 488 | if is_string(argv): 489 | argv = [argv] 490 | rv.append(Command(argv, shell=self, **opts)) 491 | return rv[0] if len(rv) == 1 else rv 492 | 493 | @contextlib.contextmanager 494 | def setenv(self, env): 495 | self.envstack.append(env) 496 | yield 497 | e = self.envstack.pop() 498 | assert e == env 499 | 500 | @contextlib.contextmanager 501 | def chdir(self, path): 502 | path = str(path) # allow pathlib.Path instances 503 | if path[0] != '/': 504 | # not absolute path, consider the current stack and join with the 505 | # last path 506 | if self.dirstack: 507 | path = os.path.normpath('{}/{}'.format(self.dirstack[-1], 508 | path)) 509 | self.dirstack.append(path) 510 | yield 511 | p = self.dirstack.pop() 512 | assert p == path 513 | 514 | def alias(self, **aliases): 515 | self.aliases.update(aliases) 516 | 517 | def export_as_module(self, module_name, full_name=False): 518 | if full_name: 519 | sys.modules[module_name] = ShellModule(self, module_name) 520 | return 521 | if module_name in globals(): 522 | raise Exception('Name "{}" is already taken'.format(module_name)) 523 | full_module_name = __name__ + '.' + module_name 524 | module = ShellModule(self, full_module_name) 525 | sys.modules[full_module_name] = module 526 | globals()[module_name] = module 527 | 528 | 529 | class ShellModule(types.ModuleType): 530 | def __init__(self, shell, name): 531 | self.__shell = shell 532 | self.__file__ = '' 533 | self.__name__ = name 534 | self.__package__ = __package__ 535 | self.__loader__ = None 536 | 537 | def __repr__(self): 538 | return repr(self.__shell) 539 | 540 | def __getattr__(self, name): 541 | if name.startswith('__'): 542 | return super(ShellModule, self).__getattr__(name) 543 | attr = getattr(self.__shell, name, None) 544 | if attr and name != 'export_as_module': 545 | return attr 546 | return self.__shell(name) 547 | 548 | 549 | class PipelineBasePy3(object): 550 | def __bytes__(self): 551 | return self._collect_output() 552 | 553 | def __str__(self): 554 | return bytes(self).decode('utf-8') 555 | 556 | 557 | class PipelineBasePy2(object): 558 | def __str__(self): 559 | return self._collect_output() 560 | 561 | def __unicode__(self): 562 | return str(self).decode('utf-8') 563 | 564 | 565 | class Pipeline(PipelineBasePy3 if PY3 else PipelineBasePy2): 566 | def __init__(self, commands): 567 | validate_pipeline(commands) 568 | self.commands = commands 569 | 570 | def __repr__(self): 571 | return ' | '.join((repr(c) for c in self.commands)) 572 | 573 | def __or__(self, other): 574 | if isinstance(other, Shell): 575 | return other(self) 576 | elif hasattr(other, 'write') or is_string(other): 577 | return Pipeline(self.commands[:-1] + 578 | [self.commands[-1]._redirect('stdout', other)]) 579 | assert isinstance(other, Command) 580 | return Pipeline(self.commands + [other]) 581 | 582 | def __ror__(self, other): 583 | if hasattr(other, '__iter__') or is_string(other): 584 | return Pipeline([self.commands[0]._redirect('stdin', other)] + 585 | self.commands[1:]) 586 | assert False, "Invalid" 587 | 588 | def __call__(self): 589 | procs, raise_on_error = self._spawn() 590 | return wait(procs, raise_on_error) 591 | 592 | def __iter__(self): 593 | return self._iter(False) 594 | 595 | def _iter(self, raw): 596 | pipeline = Pipeline(self.commands[:-1] + 597 | [self.commands[-1]._redirect('stdout', PIPE)]) 598 | procs, raise_on_error = pipeline._spawn() 599 | pipe_count = sum(1 for proc in procs if proc.stderr) 600 | if procs[-1].stdout: 601 | pipe_count += 1 602 | if not pipe_count: 603 | wait(procs, raise_on_error) 604 | # nothing to yield 605 | return 606 | iterator = iterate_outputs(procs, raise_on_error, []) 607 | if not raw: 608 | iterator = iterate_lines(iterator, trim_trailing_lf=True) 609 | if pipe_count == 1: 610 | for line, stream_index in iterator: 611 | yield line 612 | else: 613 | for line, stream_index in iterator: 614 | yield tuple(line if stream_index == index else None 615 | for index in xrange(pipe_count)) 616 | 617 | def _collect_output(self): 618 | sink = BytesIO() 619 | (self | sink)() 620 | return sink.getvalue() 621 | 622 | def _spawn(self): 623 | procs = [] 624 | raise_on_error = False 625 | for index, command in enumerate(self.commands): 626 | close_in = False 627 | close_out = False 628 | close_err = False 629 | is_first = index == 0 630 | is_last = index == len(self.commands) - 1 631 | stdin_stream = None 632 | stdout_stream = None 633 | stderr_stream = None 634 | # copy argv/opts 635 | proc_argv = [str(a) for a in command.argv] 636 | proc_opts = command.copy_opts() 637 | raise_on_error = raise_on_error or proc_opts.get('raise_on_error', 638 | False) 639 | if is_first: 640 | # first command in the pipeline may redirect stdin 641 | stdin_stream, close_in = setup_redirect(proc_opts, 'stdin') 642 | else: 643 | # only set current process stdin if it is not the first in the 644 | # pipeline. 645 | proc_opts['stdin'] = procs[-1].stdout 646 | if is_last: 647 | # last command in the pipeline may redirect stdout 648 | stdout_stream, close_out = setup_redirect(proc_opts, 'stdout') 649 | else: 650 | # only set current process stdout if it is not the last in the 651 | # pipeline. 652 | proc_opts['stdout'] = PIPE 653 | # stderr may be set at any point in the pipeline 654 | stderr_stream, close_err = setup_redirect(proc_opts, 'stderr') 655 | set_extra_popen_opts(proc_opts) 656 | if proc_opts.get('glob', False): 657 | proc_argv = expand_filenames( 658 | proc_argv, os.path.realpath( 659 | proc_opts.get('cwd', os.curdir))) 660 | set_environment(proc_opts) 661 | current_proc = RunningProcess( 662 | subprocess.Popen(proc_argv, **remove_invalid_opts(proc_opts)), 663 | stdin_stream, stdout_stream, stderr_stream, proc_argv 664 | ) 665 | # if files were opened and connected to the process stdio, close 666 | # our copies of the descriptors 667 | if close_in: 668 | proc_opts['stdin'].close() 669 | if close_out: 670 | proc_opts['stdout'].close() 671 | if close_err: 672 | proc_opts['stderr'].close() 673 | if not is_first: 674 | # close our copy of the previous process's stdout, now that it 675 | # is connected to the current process's stdin 676 | procs[-1].stdout.close() 677 | procs.append(current_proc) 678 | return procs, raise_on_error 679 | 680 | def iter_raw(self): 681 | return self._iter(True) 682 | 683 | 684 | class Command(object): 685 | OPTS = ('stdin', 'stdout', 'stderr', 'env', 'cwd', 'preexec_fn', 686 | 'raise_on_error', 'merge_env', 'glob') 687 | 688 | def __init__(self, argv, shell=None, **opts): 689 | self.argv = tuple(argv) 690 | self.shell = shell 691 | self.opts = {} 692 | for key in opts: 693 | if key not in Command.OPTS: 694 | raise TypeError('Invalid keyword argument "{}"'.format(key)) 695 | value = opts[key] 696 | if key == 'cwd' and value is not None: 697 | value = str(value) # allow pathlib.Path instances 698 | self.opts[key] = value 699 | 700 | def __call__(self, *argv, **opts): 701 | if not argv and not opts: 702 | # invoke the command 703 | return Pipeline([self])() 704 | new_opts = self.opts.copy() 705 | if 'env' in opts: 706 | update_opts_env(new_opts, opts['env']) 707 | del opts['env'] 708 | for key in Command.OPTS: 709 | if key in opts: 710 | new_opts[key] = opts[key] 711 | return Command(self.argv + argv, shell=self.shell, **new_opts) 712 | 713 | def __repr__(self): 714 | argstr = ' '.join(self.argv) 715 | optstr = ' '.join( 716 | '{}={}'.format(key, self.get_opt(key)) 717 | for key in self.iter_opts() if self.get_opt(key, None) is not None 718 | ) 719 | if optstr: 720 | return '{0} ({1})'.format(argstr, optstr) 721 | else: 722 | return argstr 723 | 724 | def __str__(self): 725 | return str(Pipeline([self])) 726 | 727 | def __bytes__(self): 728 | return bytes(Pipeline([self])) 729 | 730 | def __unicode__(self): 731 | return unicode(Pipeline([self])) 732 | 733 | def __iter__(self): 734 | return iter(Pipeline([self])) 735 | 736 | def __or__(self, other): 737 | return Pipeline([self]) | other 738 | 739 | def __ror__(self, other): 740 | return other | Pipeline([self]) 741 | 742 | def _redirect(self, key, stream): 743 | if self.get_opt(key, None) is not None: 744 | raise AlreadyRedirected('command already redirects ' + key) 745 | return self(**{key: stream}) 746 | 747 | def iter_raw(self): 748 | return Pipeline([self]).iter_raw() 749 | 750 | def get_env(self): 751 | if not self.shell.envstack and 'env' not in self.opts: 752 | return None 753 | env = {} 754 | for e in self.shell.envstack: 755 | env.update(e) 756 | env.update(self.opts.get('env', {})) 757 | return env 758 | 759 | def get_cwd(self): 760 | cwd = self.opts.get('cwd', None) 761 | if cwd: 762 | return cwd 763 | if self.shell.dirstack: 764 | return self.shell.dirstack[-1] 765 | return None 766 | 767 | def get_opt(self, opt, default=None): 768 | if opt == 'env': 769 | return self.get_env() 770 | if opt == 'cwd': 771 | return self.get_cwd() 772 | rv = self.opts.get(opt, EMPTY) 773 | if rv is EMPTY: 774 | rv = self.shell.defaults.get(opt, EMPTY) 775 | if rv is EMPTY: 776 | return default 777 | return rv 778 | 779 | def iter_opts(self): 780 | return set(list(self.opts.keys()) + list(self.shell.defaults.keys()) + 781 | ['cwd', 'env']) 782 | 783 | def copy_opts(self): 784 | rv = {} 785 | for opt in self.iter_opts(): 786 | val = self.get_opt(opt) 787 | if val is not None: 788 | rv[opt] = val 789 | return rv 790 | 791 | builtin_sh = Shell(raise_on_error=True) 792 | builtin_sh.alias( 793 | apt_cache='apt-cache', 794 | apt_get='apt-get', 795 | apt_key='apt-key', 796 | dpkg_divert='dpkg-divert', 797 | grub_install='grub-install', 798 | grub_mkconfig='grub-mkconfig', 799 | locale_gen='locale-gen', 800 | mkfs_ext2='mkfs.ext2', 801 | mkfs_ext3='mkfs.ext3', 802 | mkfs_ext4='mkfs.ext4', 803 | mkfs_vfat='mkfs.vfat', 804 | qemu_img='qemu-img', 805 | repo_add='repo-add', 806 | update_grub='update-grub', 807 | update_initramfs='update-initramfs', 808 | update_locale='update-locale', 809 | ) 810 | builtin_sh.export_as_module('sh') 811 | --------------------------------------------------------------------------------