Add some utility to IO class

cyaron/__init__.py

@@ -13,7 +13,7 @@
 from .compare import Compare
 from .consts import *
 from .graph import Edge, Graph
-from .io import IO
+from .io import IO, File
 from .math import *
 from .merger import Merger
 from .polygon import Polygon

cyaron/io.py

@@ -8,11 +8,17 @@
 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 File(IntEnum):
+    INPUT = 0
+    OUTPUT = 1
+
 
 class IO:
     """Class IO: IO tool class. It will process the input and output files."""
@@ -22,7 +28,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 +38,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 +49,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 +84,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,49 +96,56 @@ 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, File.INPUT, make_dirs)
         if not disable_output:
-            self.__init_file(output_file, data_id, "o")
+            self.__init_file(output_file, data_id, File.OUTPUT, 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: File, make_dirs: bool):
         if isinstance(f, IOBase):
             # consider ``f`` as a file object
-            if file_type == "i":
+            if file_type == File.INPUT:
                 self.input_file = f
             else:
                 self.output_file = f
         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 "")
-            if file_type == "i":
+            # be sure dir is existed
+            if make_dirs:
+                self.__make_dirs(filename)
+            if file_type == File.INPUT:
                 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 +188,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,28 +203,65 @@ def __write(self, file: IOBase, *args, **kwargs):
                 if arg == "\n":
                     self.is_first_char[file] = True
 
-    def input_write(self, *args, **kwargs):
+    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 write(self, file: File = File.INPUT, *args, **kwargs):
         """
         Write every element in *args into the input file. Splits with `separator`.
         It will convert every element into str.
         Args:
+            file: Which file (File.INPUT/File.OUTPUT) to write
             *args: the elements to write
             separator: a string used to separate every element. Defaults to " ".
         """
+        if file == File.INPUT:
+            file = self.input_file
+        elif file == File.OUTPUT:
+            file = self.output_file
+        else:
+            raise ValueError("file type is not in File.INPUT or File.OUTPUT")
+
         self.__write(self.input_file, *args, **kwargs)
 
-    def input_writeln(self, *args, **kwargs):
+    def writeln(self, file: File = File.INPUT, *args, **kwargs):
         """
         Write every element in *args into the input file and turn to a new line
         Splits with `separator`.
         It will convert every element into str.
         Args:
+            file: Which file (File.INPUT/File.OUTPUT) to write
             *args: the elements to write
             separator: a string used to separate every element. Defaults to " ".
         """
         args = list(args)
         args.append("\n")
-        self.input_write(*args, **kwargs)
+        self.write(file, *args, **kwargs)
+
+
+    def clear_content(self, file: File = File.INPUT, *args, **kwargs):
+        """
+        Clear the content of input/output
+        Args:
+            file: Which file (File.INPUT/File.OUTPUT) to clear
+            pos: Where file will truncate.
+        """
+        if file == File.INPUT:
+            file = self.input_file
+        elif file == File.OUTPUT:
+            file = self.output_file
+        else:
+            raise ValueError("file type is not in File.INPUT or File.OUTPUT")
+
+        self.__clear(file, *args, **kwargs)
 
     def output_gen(self, shell_cmd, time_limit=None):
         """
@@ -240,29 +296,6 @@ def output_gen(self, shell_cmd, time_limit=None):
 
         log.debug(self.output_filename, " done")
 
-    def output_write(self, *args, **kwargs):
-        """
-        Write every element in *args into the output file. Splits with `separator`.
-        It will convert every element into str.
-        Args:
-            *args: the elements to write
-            separator: a string used to separate every element. Defaults to " ".
-        """
-        self.__write(self.output_file, *args, **kwargs)
-
-    def output_writeln(self, *args, **kwargs):
-        """
-        Write every element in *args into the output file and turn to a new line.
-        Splits with `separator`.
-        It will convert every element into str.
-        Args:
-            *args: the elements to write
-            separator: a string used to separate every element. Defaults to " ".
-        """
-        args = list(args)
-        args.append("\n")
-        self.output_write(*args, **kwargs)
-
     def flush_buffer(self):
         """Flush the input file"""
         self.input_file.flush()