bufixes

archivebox/system.py

@@ -4,69 +4,44 @@ __package__ = 'archivebox'
 import os
 import shutil
 
-import json as pyjson
+from json import dump
+from pathlib import Path
 from typing import Optional, Union, Set, Tuple
+from subprocess import run as subprocess_run
 
 from crontab import CronTab
-from atomicwrites import atomic_write as awrite
-
-from subprocess import (
-    Popen,
-    PIPE,
-    DEVNULL, 
-    CompletedProcess,
-    TimeoutExpired,
-    CalledProcessError,
-)
+from atomicwrites import atomic_write as lib_atomic_write
 
 from .util import enforce_types, ExtendedEncoder
 from .config import OUTPUT_PERMISSIONS
 
 
 
-def run(*popenargs, input=None, capture_output=False, timeout=None, check=False, **kwargs):
+def run(*args, input=None, capture_output=True, text=True, timeout=None, check=False, **kwargs):
     """Patched of subprocess.run to fix blocking io making timeout=innefective"""
 
     if input is not None:
         if 'stdin' in kwargs:
             raise ValueError('stdin and input arguments may not both be used.')
-        kwargs['stdin'] = PIPE
 
     if capture_output:
         if ('stdout' in kwargs) or ('stderr' in kwargs):
             raise ValueError('stdout and stderr arguments may not be used '
                              'with capture_output.')
-        kwargs['stdout'] = PIPE
-        kwargs['stderr'] = PIPE
 
-    with Popen(*popenargs, **kwargs) as process:
-        try:
-            stdout, stderr = process.communicate(input, timeout=timeout)
-        except TimeoutExpired:
-            process.kill()
-            try:
-                stdout, stderr = process.communicate(input, timeout=2)
-            except:
-                pass
-            raise TimeoutExpired(popenargs[0][0], timeout)
-        except BaseException:
-            process.kill()
-            # We don't call process.wait() as .__exit__ does that for us.
-            raise 
-        retcode = process.poll()
-        if check and retcode:
-            raise CalledProcessError(retcode, process.args,
-                                     output=stdout, stderr=stderr)
-    return CompletedProcess(process.args, retcode, stdout, stderr)
+    return subprocess_run(*args, input=input, capture_output=capture_output, text=text, timeout=timeout, check=check, **kwargs)
 
-
-def atomic_write(path: str, contents: Union[dict, str, bytes], overwrite: bool=True) -> None:
+@enforce_types
+def atomic_write(path: Union[Path, str], contents: Union[dict, str, bytes], overwrite: bool=True) -> None:
     """Safe atomic write to filesystem by writing to temp file + atomic rename"""
     
-    with awrite(path, overwrite=overwrite) as f:
+    mode = 'wb+' if isinstance(contents, bytes) else 'w'
+
+    # print('\n> Atomic Write:', mode, path, len(contents), f'overwrite={overwrite}')
+    with lib_atomic_write(path, mode=mode, overwrite=overwrite) as f:
         if isinstance(contents, dict):
-            pyjson.dump(contents, f, indent=4, sort_keys=True, cls=ExtendedEncoder)
-        else:
+            dump(contents, f, indent=4, sort_keys=True, cls=ExtendedEncoder)
+        elif isinstance(contents, (bytes, str)):
             f.write(contents)
 
 @enforce_types
@@ -76,7 +51,7 @@ def chmod_file(path: str, cwd: str='.', permissions: str=OUTPUT_PERMISSIONS, tim
     if not os.path.exists(os.path.join(cwd, path)):
         raise Exception('Failed to chmod: {} does not exist (did the previous step fail?)'.format(path))
 
-    chmod_result = run(['chmod', '-R', permissions, path], cwd=cwd, stdout=DEVNULL, stderr=PIPE, timeout=timeout)
+    chmod_result = run(['chmod', '-R', permissions, path], cwd=cwd, timeout=timeout)
     if chmod_result.returncode == 1:
         print('     ', chmod_result.stderr.decode())
         raise Exception('Failed to chmod {}/{}'.format(cwd, path))