wmaee.core.io

  1from os import getcwd, mkdir, chdir
  2from os.path import exists, split
  3from uuid import uuid4
  4from shutil import rmtree
  5from typing import Optional, Generator, List
  6
  7class working_directory:
  8    """
  9    A convenience class which provides syntactic sugar, allowing the user to change directories.
 10    Can also be nested.
 11    
 12    Parameters
 13    ----------
 14    name : Optional[str], optional
 15        Name of the directory. If None is given, `os.getcwd()` will be used. Default is None.
 16    create : bool, optional
 17        Whether to create the directory if it doesn't exist. Default is True.
 18    delete : bool, optional
 19        Whether to delete the directory after a 'with' clause. Default is False.
 20    """
 21    def __init__(self, name: Optional[str] = None, create: bool = True, delete: bool = False):
 22        self._name: str = str(uuid4()) if not name else name
 23        self._delete: bool = delete
 24        self._create: bool = create
 25        self._curr_dir: str = getcwd()
 26        self._active: bool = False
 27        
 28    def _split_path(self, path: str) -> List[str]:
 29        """
 30        Helper function to split a string into a list of directories.
 31        
 32        Parameters
 33        ----------
 34        path : str
 35            The path to be split.
 36
 37        Returns
 38        -------
 39        list[str]
 40            List of directories.
 41        """
 42        allparts: List[str] = []
 43        while 1:
 44            parts = split(path)
 45            if parts[0] == path:  # Sentinel for absolute paths
 46                allparts.insert(0, parts[0])
 47                break
 48            elif parts[1] == path:  # Sentinel for relative paths
 49                allparts.insert(0, parts[1])
 50                break
 51            else:
 52                path = parts[0]
 53                allparts.insert(0, parts[1])
 54        return allparts
 55
 56    def __enter__(self) -> 'working_directory':
 57        """
 58        Enter the working directory.
 59
 60        Returns
 61        -------
 62        working_directory
 63            The working_directory object.
 64        """
 65        path = self._split_path(self._name)
 66        for sub in path:
 67            if not exists(sub):
 68                if self._create:
 69                    mkdir(sub)
 70                else:
 71                    raise Exception(f"The requested directory {sub} in {getcwd()} "
 72                                    f"doesn't exist, and you do not want me to create it.")
 73            chdir(sub)
 74        self._active = True
 75        return self
 76
 77    def __exit__(self, exc_type, exc_val, exc_tb):
 78        """
 79        Exit the working directory.
 80
 81        Parameters
 82        ----------
 83        exc_type : type
 84            Type of the exception.
 85        exc_val : any
 86            The exception value.
 87        exc_tb : traceback
 88            The traceback object.
 89        """
 90        chdir(self._curr_dir)
 91        if self._delete:
 92            rmtree(self._name)
 93        self._active = False
 94
 95    @property
 96    def name(self) -> str:
 97        """
 98        Get the name of the working directory.
 99
100        Returns
101        -------
102        str
103            The name of the working directory.
104        """
105        return self._name
106
107    @property
108    def active(self) -> bool:
109        """
110        Check if the working directory is active.
111
112        Returns
113        -------
114        bool
115            True if the working directory is active, False otherwise.
116        """
117        return self._active
118
119
120# Utility functions
121def grep(file: str, string: str) -> Generator[str, None, None]:
122    """
123    A simple Python implementation of the Linux `grep` command.
124    It goes through a (text) file and returns all lines containing a desired string as a generator.
125
126    Parameters
127    ----------
128    file : str
129        Name of the file to be opened for reading and searching of the string.
130    string : str
131        String to be searched for.
132
133    Yields
134    ------
135    str
136        (Next) line containing the string `string` in the file `file`.
137    """    
138    with open(file, 'r') as f:
139        for line in f:
140            if string in line:
141                yield line