File size: 4,616 Bytes
c8e7ce2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 |
#!/usr/bin/env python
# coding=utf-8
# Copyright 2021 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License
"""Contains utilities to easily handle subprocesses in `huggingface_hub`."""
import os
import subprocess
import sys
from contextlib import contextmanager
from io import StringIO
from pathlib import Path
from typing import IO, Generator, List, Optional, Tuple, Union
from .logging import get_logger
logger = get_logger(__name__)
@contextmanager
def capture_output() -> Generator[StringIO, None, None]:
"""Capture output that is printed to terminal.
Taken from https://stackoverflow.com/a/34738440
Example:
```py
>>> with capture_output() as output:
... print("hello world")
>>> assert output.getvalue() == "hello world\n"
```
"""
output = StringIO()
previous_output = sys.stdout
sys.stdout = output
yield output
sys.stdout = previous_output
def run_subprocess(
command: Union[str, List[str]],
folder: Optional[Union[str, Path]] = None,
check=True,
**kwargs,
) -> subprocess.CompletedProcess:
"""
Method to run subprocesses. Calling this will capture the `stderr` and `stdout`,
please call `subprocess.run` manually in case you would like for them not to
be captured.
Args:
command (`str` or `List[str]`):
The command to execute as a string or list of strings.
folder (`str`, *optional*):
The folder in which to run the command. Defaults to current working
directory (from `os.getcwd()`).
check (`bool`, *optional*, defaults to `True`):
Setting `check` to `True` will raise a `subprocess.CalledProcessError`
when the subprocess has a non-zero exit code.
kwargs (`Dict[str]`):
Keyword arguments to be passed to the `subprocess.run` underlying command.
Returns:
`subprocess.CompletedProcess`: The completed process.
"""
if isinstance(command, str):
command = command.split()
if isinstance(folder, Path):
folder = str(folder)
return subprocess.run(
command,
stderr=subprocess.PIPE,
stdout=subprocess.PIPE,
check=check,
encoding="utf-8",
errors="replace", # if not utf-8, replace char by �
cwd=folder or os.getcwd(),
**kwargs,
)
@contextmanager
def run_interactive_subprocess(
command: Union[str, List[str]],
folder: Optional[Union[str, Path]] = None,
**kwargs,
) -> Generator[Tuple[IO[str], IO[str]], None, None]:
"""Run a subprocess in an interactive mode in a context manager.
Args:
command (`str` or `List[str]`):
The command to execute as a string or list of strings.
folder (`str`, *optional*):
The folder in which to run the command. Defaults to current working
directory (from `os.getcwd()`).
kwargs (`Dict[str]`):
Keyword arguments to be passed to the `subprocess.run` underlying command.
Returns:
`Tuple[IO[str], IO[str]]`: A tuple with `stdin` and `stdout` to interact
with the process (input and output are utf-8 encoded).
Example:
```python
with _interactive_subprocess("git credential-store get") as (stdin, stdout):
# Write to stdin
stdin.write("url=hf.co\nusername=obama\n".encode("utf-8"))
stdin.flush()
# Read from stdout
output = stdout.read().decode("utf-8")
```
"""
if isinstance(command, str):
command = command.split()
with subprocess.Popen(
command,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
encoding="utf-8",
errors="replace", # if not utf-8, replace char by �
cwd=folder or os.getcwd(),
**kwargs,
) as process:
assert process.stdin is not None, "subprocess is opened as subprocess.PIPE"
assert process.stdout is not None, "subprocess is opened as subprocess.PIPE"
yield process.stdin, process.stdout
|