Add some utility to IO class

cyaron/io.py

@@ -8,12 +8,13 @@
 import re
 import subprocess
 import tempfile
+from enum import IntEnum
+from pathlib import Path
 from typing import Union, overload
 from io import IOBase
 from . import log
 from .utils import list_like, make_unicode
 
-
 class IO:
     """Class IO: IO tool class. It will process the input and output files."""
 
@@ -22,7 +23,8 @@ def __init__(self,
                  input_file: Union[IOBase, str, int, None] = None,
                  output_file: Union[IOBase, str, int, None] = None,
                  data_id: Union[str, None] = None,
-                 disable_output: bool = False):
+                 disable_output: bool = False,
+                 make_dirs: bool = False):
         ...
 
     @overload
@@ -31,7 +33,8 @@ def __init__(self,
                  file_prefix: Union[str, None] = None,
                  input_suffix: Union[str, None] = '.in',
                  output_suffix: Union[str, None] = '.out',
-                 disable_output: bool = False):
+                 disable_output: bool = False,
+                 make_dirs: bool = False):
         ...
 
     def __init__(self,
@@ -41,20 +44,22 @@ def __init__(self,
                  file_prefix: Union[str, None] = None,
                  input_suffix: Union[str, None] = '.in',
                  output_suffix: Union[str, None] = '.out',
-                 disable_output: bool = False):
+                 disable_output: bool = False,
+                 make_dirs: bool = False):
         """
         Args:
-            input_file (optional): input file object or filename or file descriptor. 
+            input_file (optional): input file object or filename or file descriptor.
                 If it's None, make a temp file. Defaults to None.
-            output_file (optional): input file object or filename or file descriptor. 
+            output_file (optional): input file object or filename or file descriptor.
                 If it's None, make a temp file. Defaults to None.
-            data_id (optional): the id of the data. It will be add after 
+            data_id (optional): the id of the data. It will be add after
                 `input_file` and `output_file` when they are str.
                 If it's None, the file names will not contain the id. Defaults to None.
             file_prefix (optional): the prefix for the input and output files. Defaults to None.
             input_suffix (optional): the suffix of the input file. Defaults to '.in'.
             output_suffix (optional): the suffix of the output file. Defaults to '.out'.
             disable_output (optional): set to True to disable output file. Defaults to False.
+            make_dirs (optional): set to True to create dir if path is not found. Defaults to True.
         Examples:
             >>> IO("a","b")
             # create input file "a" and output file "b"
@@ -74,6 +79,8 @@ def __init__(self,
             # create input file "data2.in" and output file "data2.out"
             >>> IO(open('data.in', 'w+'), open('data.out', 'w+'))
             # input file "data.in" and output file "data.out"
+            >>> IO("./io/data.in", "./io/data.out")
+            # input file "./io/data.in" and output file "./io/data.out" if the dir "./io" not found it will be created
         """
         if file_prefix is not None:
             # legacy mode
@@ -84,16 +91,16 @@ def __init__(self,
                 self.__escape_format(output_suffix))
         self.input_filename, self.output_filename = None, None
         self.__input_temp, self.__output_temp = False, False
-        self.__init_file(input_file, data_id, "i")
+        self.__init_file(input_file, data_id, "i", make_dirs)
         if not disable_output:
-            self.__init_file(output_file, data_id, "o")
+            self.__init_file(output_file, data_id, "o", make_dirs)
         else:
             self.output_file = None
         self.__closed = False
         self.is_first_char = {}
 
     def __init_file(self, f: Union[IOBase, str, int, None],
-                    data_id: Union[int, None], file_type: str):
+                    data_id: Union[int, None], file_type: str, make_dirs: bool):
         if isinstance(f, IOBase):
             # consider ``f`` as a file object
             if file_type == "i":
@@ -103,30 +110,37 @@ def __init_file(self, f: Union[IOBase, str, int, None],
         elif isinstance(f, int):
             # consider ``f`` as a file descor
             self.__init_file(open(f, 'w+', encoding="utf-8", newline='\n'),
-                             data_id, file_type)
+                             data_id, file_type, make_dirs)
         elif f is None:
             # consider wanna temp file
             fd, self.input_filename = tempfile.mkstemp()
-            self.__init_file(fd, data_id, file_type)
-            if file_type == "i":
+            self.__init_file(fd, data_id, file_type, make_dirs)
+            if file_type == File.INPUT:
                 self.__input_temp = True
             else:
                 self.__output_temp = True
         else:
             # consider ``f`` as filename template
             filename = f.format(data_id or "")
+            # be sure dir is existed
+            if make_dirs:
+                self.__make_dirs(filename)
             if file_type == "i":
                 self.input_filename = filename
             else:
                 self.output_filename = filename
             self.__init_file(
                 open(filename, 'w+', newline='\n', encoding='utf-8'), data_id,
-                file_type)
+                file_type, make_dirs)
 
     def __escape_format(self, st: str):
         """replace "{}" to "{{}}" """
         return re.sub(r"\{", "{{", re.sub(r"\}", "}}", st))
 
+    def __make_dirs(self, pth: Union[IOBase, str, int, None]):
+        if isinstance(pth, str):
+            os.makedirs(os.path.dirname(pth), exist_ok=True)
+
     def __del_files(self):
         """delete files"""
         if self.__input_temp and self.input_filename is not None:
@@ -169,7 +183,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     def __write(self, file: IOBase, *args, **kwargs):
         """
-        Write every element in *args into file. If the element isn't "\n", insert `separator`. 
+        Write every element in *args into file. If the element isn't "\n", insert `separator`.
         It will convert every element into str.
         """
         separator = kwargs.get("separator", " ")
@@ -184,6 +198,17 @@ def __write(self, file: IOBase, *args, **kwargs):
                 if arg == "\n":
                     self.is_first_char[file] = True
 
+    def __clear(self, file: IOBase, *args, **kwargs):
+        """
+        Clear the content use truncate()
+        Args:
+            file: Which file (File.INPUT/File.OUTPUT) to clear
+            pos: Where file will truncate.
+        """
+        pos = kwargs.get("pos", 0)
+        file.truncate(pos)
+
+
     def input_write(self, *args, **kwargs):
         """
         Write every element in *args into the input file. Splits with `separator`.
@@ -207,6 +232,15 @@ def input_writeln(self, *args, **kwargs):
         args.append("\n")
         self.input_write(*args, **kwargs)
 
+    def input_clear_content(self, *args, **kwargs):
+        """
+        Clear the content of input
+        Args:
+            pos: Where file will truncate.
+        """
+
+        self.__clear(self.input_file, *args, **kwargs)
+
     def output_gen(self, shell_cmd, time_limit=None):
         """
         Run the command `shell_cmd` (usually the std program) and send it the input file as stdin.
@@ -263,6 +297,15 @@ def output_writeln(self, *args, **kwargs):
         args.append("\n")
         self.output_write(*args, **kwargs)
 
+
+    def output_clear_content(self, *args, **kwargs):
+        """
+        Clear the content of output
+        Args:
+            pos: Where file will truncate
+        """
+        self.__clear(self.output_file, *args, **kwargs)
+
     def flush_buffer(self):
         """Flush the input file"""
         self.input_file.flush()