megadetector 5.0.6__py3-none-any.whl → 5.0.8__py3-none-any.whl

md_utils/path_utils.py

@@ -12,16 +12,23 @@
 import glob
 import ntpath
 import os
+import sys
+import platform
 import posixpath
 import string
 import json
+import shutil
 import unicodedata
 import zipfile
+import webbrowser
+import subprocess
+import re
 
 from zipfile import ZipFile
 from datetime import datetime
 from typing import Container, Iterable, List, Optional, Tuple, Sequence
-from multiprocessing.pool import ThreadPool
+from multiprocessing.pool import Pool, ThreadPool
+from functools import partial
 from tqdm import tqdm
 
 IMG_EXTENSIONS = ('.jpg', '.jpeg', '.gif', '.png', '.tif', '.tiff', '.bmp')
@@ -34,31 +41,53 @@ CHAR_LIMIT = 255
 
 #%% General path functions
 
-def recursive_file_list(base_dir, convert_slashes=True, return_relative_paths=False):
-    """
+def recursive_file_list(base_dir, convert_slashes=True, 
+                        return_relative_paths=False, sort_files=True,
+                        recursive=True):
+    r"""
     Enumerate files (not directories) in [base_dir], optionally converting
     \ to /
     """
     
+    assert os.path.isdir(base_dir), '{} is not a folder'.format(base_dir)
+    
     all_files = []
 
-    for root, _, filenames in os.walk(base_dir):
-        for filename in filenames:
-            full_path = os.path.join(root, filename)
-            all_files.append(full_path)
-
+    if recursive:
+        for root, _, filenames in os.walk(base_dir):
+            for filename in filenames:
+                full_path = os.path.join(root, filename)
+                all_files.append(full_path)
+    else:
+        all_files_relative = os.listdir(base_dir)
+        all_files = [os.path.join(base_dir,fn) for fn in all_files_relative]
+        all_files = [fn for fn in all_files if os.path.isfile(fn)]
+        
     if return_relative_paths:
         all_files = [os.path.relpath(fn,base_dir) for fn in all_files]
 
     if convert_slashes:
         all_files = [fn.replace('\\', '/') for fn in all_files]
+    
+    if sort_files:
+        all_files = sorted(all_files)
         
-    all_files = sorted(all_files)
     return all_files
 
 
-def split_path(path: str) -> List[str]:
+def file_list(base_dir, convert_slashes=True, return_relative_paths=False, sort_files=True, 
+              recursive=False):
     """
+    Trivial wrapper for recursive_file_list, which was a poor function name choice at the time, 
+    it doesn't really make sense to have a "recursive" option in a function called "recursive_file_list".
+    """
+    
+    return recursive_file_list(base_dir,convert_slashes,return_relative_paths,sort_files,
+                               recursive=recursive)
+
+
+def split_path(path: str) -> List[str]:
+    r"""
     Splits [path] into all its constituent tokens.
 
     Non-recursive version of:
@@ -88,7 +117,7 @@ def split_path(path: str) -> List[str]:
 
 
 def fileparts(path: str) -> Tuple[str, str, str]:
-    """
+    r"""
     Breaks down a path into the directory path, filename, and extension.
 
     Note that the '.' lives with the extension, and separators are removed.
@@ -187,7 +216,8 @@ def safe_create_link(link_exists,link_new):
     it.
     
     Errors if link_new already exists but it's not a link.
-    """    
+    """
+    
     if os.path.exists(link_new) or os.path.islink(link_new):
         assert os.path.islink(link_new)
         if not os.readlink(link_new) == link_exists:
@@ -197,23 +227,6 @@ def safe_create_link(link_exists,link_new):
         os.symlink(link_exists,link_new)
         
 
-def get_file_sizes(base_dir, convert_slashes=True):
-    """
-    Get sizes recursively for all files in base_dir, returning a dict mapping
-    relative filenames to size.
-    """
-    
-    relative_filenames = recursive_file_list(base_dir, convert_slashes=convert_slashes, 
-                                             return_relative_paths=True)
-    
-    fn_to_size = {}
-    for fn_relative in tqdm(relative_filenames):
-        fn_abs = os.path.join(base_dir,fn_relative)
-        fn_to_size[fn_relative] = os.path.getsize(fn_abs)
-                   
-    return fn_to_size
-        
-
 #%% Image-related path functions
 
 def is_image_file(s: str, img_extensions: Container[str] = IMG_EXTENSIONS
@@ -240,14 +253,17 @@ def find_image_strings(strings: Iterable[str]) -> List[str]:
 
 
 def find_images(dirname: str, recursive: bool = False, 
-                return_relative_paths: bool = False, convert_slashes: bool = False) -> List[str]:
+                return_relative_paths: bool = False, 
+                convert_slashes: bool = False) -> List[str]:
     """
     Finds all files in a directory that look like image file names. Returns
     absolute paths unless return_relative_paths is set.  Uses the OS-native
-    path separator unless convert_slahes is set, in which case will always
+    path separator unless convert_slashes is set, in which case will always
     use '/'.
     """
     
+    assert os.path.isdir(dirname), '{} is not a folder'.format(dirname)
+    
     if recursive:
         strings = glob.glob(os.path.join(dirname, '**', '*.*'), recursive=True)
     else:
@@ -270,11 +286,11 @@ def find_images(dirname: str, recursive: bool = False,
 
 def clean_filename(filename: str, allow_list: str = VALID_FILENAME_CHARS,
                    char_limit: int = CHAR_LIMIT, force_lower: bool = False) -> str:
-    """
+    r"""
     Removes non-ASCII and other invalid filename characters (on any
     reasonable OS) from a filename, then trims to a maximum length.
 
-    Does not allow :\/, use clean_path if you want to preserve those.
+    Does not allow :\/ by default, use clean_path if you want to preserve those.
 
     Adapted from
     https://gist.github.com/wassname/1393c4a57cfcbf03641dbc31886123b8
@@ -319,15 +335,91 @@ def flatten_path(pathname: str, separator_chars: str = SEPARATOR_CHARS) -> str:
 
 #%% Platform-independent way to open files in their associated application
 
-import sys,subprocess
+def environment_is_wsl():
+    """
+    Returns True if we're running in WSL
+    """
+    
+    if sys.platform not in ('linux','posix'):
+        return False
+    platform_string = ' '.join(platform.uname()).lower()
+    return 'microsoft' in platform_string and 'wsl' in platform_string
+    
 
-def open_file(filename):
-    if sys.platform == "win32":
+def wsl_path_to_windows_path(filename):
+    """
+    Converts a WSL path to a Windows path, or returns None if that's not possible.  E.g.
+    converts:
+        
+    /mnt/e/a/b/c
+    
+    ...to:
+        
+    e:\a\b\c
+    """
+    
+    result = subprocess.run(['wslpath', '-w', filename], text=True, capture_output=True)
+    if result.returncode != 0:
+        print('Could not convert path {} from WSL to Windows'.format(filename))
+        return None
+    return result.stdout.strip()
+    
+    
+def open_file(filename, attempt_to_open_in_wsl_host=False, browser_name=None):
+    """
+    Opens [filename] in the default OS file handler for this file type.
+    
+    If attempt_to_open_in_wsl_host is True, and we're in WSL, attempts to open
+    [filename] in the Windows host environment.
+    
+    If browser_name is not None, uses the webbrowser module to open the filename
+    in the specified browser; see https://docs.python.org/3/library/webbrowser.html
+    for supported browsers.  Falls back to the default file handler if webbrowser.open()
+    fails.  In this case, attempt_to_open_in_wsl_host is ignored unless webbrowser.open() fails.
+    
+    If browser_name is 'default', use the system default.  This is different from the 
+    parameter to webbrowser.get(), where None implies the system default.
+    """
+    
+    if browser_name is not None:
+        if browser_name == 'chrome':
+            browser_name = 'google-chrome'
+        elif browser_name == 'default':
+            browser_name = None
+        try:
+            result = webbrowser.get(using=browser_name).open(filename)
+        except Exception:
+            result = False
+        if result:
+            return
+        
+    if sys.platform == 'win32':
+        
         os.startfile(filename)
+
+    elif sys.platform == 'darwin':
+      
+        opener = 'open'
+        subprocess.call([opener, filename])
+            
+    elif attempt_to_open_in_wsl_host and environment_is_wsl():
+        
+        windows_path = wsl_path_to_windows_path(filename)
+        
+        # Fall back to xdg-open
+        if windows_path is None:
+            subprocess.call(['xdg-open', filename])
+            
+        if os.path.isdir(filename):            
+            subprocess.run(["explorer.exe", windows_path])
+        else:
+            os.system("cmd.exe /C start %s" % (re.escape(windows_path)))    
+        
     else:
-        opener = "open" if sys.platform == "darwin" else "xdg-open"
+        
+        opener = 'xdg-open'        
         subprocess.call([opener, filename])
-
+        
 
 #%% File list functions
 
@@ -358,6 +450,107 @@ def read_list_from_file(filename: str) -> List[str]:
     return file_list
 
 
+def _copy_file(input_output_tuple,overwrite=True,verbose=False):
+    assert len(input_output_tuple) == 2
+    source_fn = input_output_tuple[0]
+    target_fn = input_output_tuple[1]
+    if (not overwrite) and (os.path.isfile(target_fn)):
+        if verbose:
+            print('Skipping existing file {}'.format(target_fn))
+        return
+    os.makedirs(os.path.dirname(target_fn),exist_ok=True)
+    shutil.copyfile(source_fn,target_fn)
+    
+
+def parallel_copy_files(input_file_to_output_file, max_workers=16, 
+                        use_threads=True, overwrite=False, verbose=False):
+    """
+    Copy files from source to target according to the dict input_file_to_output_file.
+    """
+
+    n_workers = min(max_workers,len(input_file_to_output_file))
+    
+    # Package the dictionary as a set of 2-tuples
+    input_output_tuples = []
+    for input_fn in input_file_to_output_file:
+        input_output_tuples.append((input_fn,input_file_to_output_file[input_fn]))
+
+    if use_threads:
+        pool = ThreadPool(n_workers)
+    else:
+        pool = Pool(n_workers)
+
+    with tqdm(total=len(input_output_tuples)) as pbar:
+        for i,_ in enumerate(pool.imap_unordered(partial(_copy_file,overwrite=overwrite,verbose=verbose),
+                                                 input_output_tuples)):
+            pbar.update()
+
+# ...def parallel_copy_files(...)
+
+
+def get_file_sizes(base_dir, convert_slashes=True):
+    """
+    Get sizes recursively for all files in base_dir, returning a dict mapping
+    relative filenames to size.
+    
+    TODO: merge the functionality here with parallel_get_file_sizes, which uses slightly
+    different semantics.
+    """
+    
+    relative_filenames = recursive_file_list(base_dir, convert_slashes=convert_slashes, 
+                                             return_relative_paths=True)
+    
+    fn_to_size = {}
+    for fn_relative in tqdm(relative_filenames):
+        fn_abs = os.path.join(base_dir,fn_relative)
+        fn_to_size[fn_relative] = os.path.getsize(fn_abs)
+                   
+    return fn_to_size
+        
+
+def _get_file_size(filename,verbose=False):
+    """
+    Internal function for safely getting the size of a file.  Returns a (filename,size)
+    tuple, where size is None if there is an error.
+    """
+    
+    try:
+        size = os.path.getsize(filename)
+    except Exception as e:
+        if verbose:
+            print('Error reading file size for {}: {}'.format(filename,str(e)))
+        size = None
+    return (filename,size)
+
+    
+def parallel_get_file_sizes(filenames, max_workers=16, 
+                        use_threads=True, verbose=False,
+                        recursive=True):
+    """
+    Return a dictionary mapping every file in [filenames] to the corresponding file size,
+    or None for errors.  If [filenames] is a folder, will enumerate the folder (optionally recursively).
+    """
+
+    n_workers = min(max_workers,len(filenames))
+    
+    if isinstance(filenames,str) and os.path.isdir(filenames):
+        filenames = recursive_file_list(filenames,recursive=recursive,return_relative_paths=False)
+    
+    if use_threads:
+        pool = ThreadPool(n_workers)
+    else:
+        pool = Pool(n_workers)
+
+    resize_results = list(tqdm(pool.imap(
+        partial(_get_file_size,verbose=verbose),filenames), total=len(filenames)))
+    
+    to_return = {}
+    for r in resize_results:
+        to_return[r[0]] = r[1]
+
+    return to_return
+
+
 #%% Zip functions
 
 def zip_file(input_fn, output_fn=None, overwrite=False, verbose=False, compresslevel=9):
@@ -375,7 +568,7 @@ def zip_file(input_fn, output_fn=None, overwrite=False, verbose=False, compressl
         return
     
     if verbose:
-        print('Zipping {} to {}'.format(input_fn,output_fn))
+        print('Zipping {} to {} with level {}'.format(input_fn,output_fn,compresslevel))
     
     with ZipFile(output_fn,'w',zipfile.ZIP_DEFLATED) as zipf:
         zipf.write(input_fn,arcname=basename,compresslevel=compresslevel,
@@ -384,9 +577,37 @@ def zip_file(input_fn, output_fn=None, overwrite=False, verbose=False, compressl
     return output_fn
 
 
+def zip_files_into_single_zipfile(input_files, output_fn, arc_name_base,
+                                  overwrite=False, verbose=False, compresslevel=9):
+    """
+    Zip all the files in [input_files] into [output_fn].  Archive names are relative to 
+    arc_name_base.
+    """
+    
+    if not overwrite:
+        if os.path.isfile(output_fn):
+            print('Zip file {} exists, skipping'.format(output_fn))
+            return            
+        
+    if verbose:
+        print('Zipping {} files to {} (compression level {})'.format(
+            len(input_files),output_fn,compresslevel))
+        
+    with ZipFile(output_fn,'w',zipfile.ZIP_DEFLATED) as zipf:
+        for input_fn_abs in tqdm(input_files,disable=(not verbose)):
+            input_fn_relative = os.path.relpath(input_fn_abs,arc_name_base)
+            zipf.write(input_fn_abs,
+                       arcname=input_fn_relative,
+                       compresslevel=compresslevel,
+                       compress_type=zipfile.ZIP_DEFLATED)
+
+    return output_fn
+    
+    
 def zip_folder(input_folder, output_fn=None, overwrite=False, verbose=False, compresslevel=9):
     """
-    Recursively zip everything in [input_folder], storing outputs as relative paths.
+    Recursively zip everything in [input_folder] into a single zipfile, storing outputs as relative 
+    paths.
     
     Defaults to writing to [input_folder].zip
     """
@@ -395,15 +616,18 @@ def zip_folder(input_folder, output_fn=None, overwrite=False, verbose=False, com
         output_fn = input_folder + '.zip'
         
     if not overwrite:
-        assert not os.path.isfile(output_fn), 'Zip file {} exists'.format(output_fn)
+        if os.path.isfile(output_fn):
+            print('Zip file {} exists, skipping'.format(output_fn))
+            return            
         
     if verbose:
-        print('Zipping {} to {}'.format(input_folder,output_fn))
+        print('Zipping {} to {} (compression level {})'.format(
+            input_folder,output_fn,compresslevel))
     
     relative_filenames = recursive_file_list(input_folder,return_relative_paths=True)
     
     with ZipFile(output_fn,'w',zipfile.ZIP_DEFLATED) as zipf:
-        for input_fn_relative in relative_filenames:
+        for input_fn_relative in tqdm(relative_filenames,disable=(not verbose)):
             input_fn_abs = os.path.join(input_folder,input_fn_relative)            
             zipf.write(input_fn_abs,
                        arcname=input_fn_relative,
@@ -413,19 +637,74 @@ def zip_folder(input_folder, output_fn=None, overwrite=False, verbose=False, com
     return output_fn
 
         
-def parallel_zip_files(input_files,max_workers=16):
+def parallel_zip_files(input_files, max_workers=16, use_threads=True, compresslevel=9, 
+                       overwrite=False, verbose=False):
     """
     Zip one or more files to separate output files in parallel, leaving the 
-    original files in place.
+    original files in place.  Each file is zipped to [filename].zip.
     """
 
     n_workers = min(max_workers,len(input_files))
-    pool = ThreadPool(n_workers)
+
+    if use_threads:
+        pool = ThreadPool(n_workers)
+    else:
+        pool = Pool(n_workers)
+
     with tqdm(total=len(input_files)) as pbar:
-        for i,_ in enumerate(pool.imap_unordered(zip_file,input_files)):
+        for i,_ in enumerate(pool.imap_unordered(partial(zip_file,
+          output_fn=None,overwrite=overwrite,verbose=verbose,compresslevel=compresslevel),
+          input_files)):
             pbar.update()
 
 
+def parallel_zip_folders(input_folders, max_workers=16, use_threads=True,
+                         compresslevel=9, overwrite=False, verbose=False):
+    """
+    Zip one or more folders to separate output files in parallel, leaving the 
+    original folders in place.  Each folder is zipped to [folder_name].zip.
+    """
+
+    n_workers = min(max_workers,len(input_folders))
+
+    if use_threads:
+        pool = ThreadPool(n_workers)
+    else:
+        pool = Pool(n_workers)
+    
+    with tqdm(total=len(input_folders)) as pbar:
+        for i,_ in enumerate(pool.imap_unordered(
+                partial(zip_folder,overwrite=overwrite,
+                        compresslevel=compresslevel,verbose=verbose),
+                input_folders)):
+            pbar.update()
+
+
+def zip_each_file_in_folder(folder_name,recursive=False,max_workers=16,use_threads=True,
+                            compresslevel=9,overwrite=False,required_token=None,verbose=False,
+                            exclude_zip=True):
+    """
+    Zip each file in [folder_name] to its own zipfile (filename.zip), optionally recursing.  To zip a whole
+    folder into a single zipfile, use zip_folder().
+    
+    If required_token is not None, include only files that contain that token.
+    """
+    
+    assert os.path.isdir(folder_name), '{} is not a folder'.format(folder_name)
+    
+    input_files = recursive_file_list(folder_name,recursive=recursive,return_relative_paths=False)
+    
+    if required_token is not None:
+        input_files = [fn for fn in input_files if required_token in fn]
+    
+    if exclude_zip:
+        input_files = [fn for fn in input_files if (not fn.endswith('.zip'))]
+                                                    
+    parallel_zip_files(input_files=input_files,max_workers=max_workers,
+                       use_threads=use_threads,compresslevel=compresslevel,
+                       overwrite=overwrite,verbose=verbose)
+
+
 def unzip_file(input_file, output_folder=None):
     """
     Unzip a zipfile to the specified output folder, defaulting to the same location as

md_utils/process_utils.py

@@ -17,14 +17,28 @@ import subprocess
 
 os.environ["PYTHONUNBUFFERED"] = "1"
 
-def execute(cmd):
+def execute(cmd,encoding=None,errors=None,env=None,verbose=False):
     """
     Run [cmd] (a single string) in a shell, yielding each line of output to the caller.
+    
+    The "encoding", "errors", and "env" parameters are passed directly to subprocess.Popen().
+    
+    "verbose" only impacts output about process management, it is not related to printing
+    output from the child process.
     """
- 
+    
+    if verbose: 
+        if encoding is not None:
+            print('Launching child process with non-default encoding {}'.format(encoding))
+        if errors is not None:
+            print('Launching child process with non-default text error handling {}'.format(errors))
+        if env is not None:
+            print('Launching child process with non-default environment {}'.format(str(env)))
+        
     # https://stackoverflow.com/questions/4417546/constantly-print-subprocess-output-while-process-is-running
     popen = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
-                             shell=True, universal_newlines=True)
+                             shell=True, universal_newlines=True, encoding=encoding,
+                             errors=errors, env=env)
     for stdout_line in iter(popen.stdout.readline, ""):
         yield stdout_line
     popen.stdout.close()
@@ -33,22 +47,27 @@ def execute(cmd):
         raise subprocess.CalledProcessError(return_code, cmd)
 
 
-def execute_and_print(cmd,print_output=True):
+def execute_and_print(cmd,print_output=True,encoding=None,errors=None,env=None,verbose=False):
     """
     Run [cmd] (a single string) in a shell, capturing and printing output.  Returns
     a dictionary with fields "status" and "output".
+    
+    The "encoding", "errors", and "env" parameters are passed directly to subprocess.Popen().
+    
+    "verbose" only impacts output about process management, it is not related to printing
+    output from the child process.
     """
 
     to_return = {'status':'unknown','output':''}
-    output=[]
+    output = []
     try:
-        for s in execute(cmd):
+        for s in execute(cmd,encoding=encoding,errors=errors,env=env,verbose=verbose):
             output.append(s)
             if print_output:
                 print(s,end='',flush=True)
         to_return['status'] = 0
     except subprocess.CalledProcessError as cpe:
-        print('execute_and_print caught error: {}'.format(cpe.output))
+        print('execute_and_print caught error: {} ({})'.format(cpe.output,str(cpe)))
         to_return['status'] = cpe.returncode
     to_return['output'] = output