refactor: refactor writers & processor utils

Author: boomb0om

DPF/processors/helpers/dataframe_helper.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Any
 
 import pandas as pd
 from tqdm.contrib.concurrent import thread_map
@@ -19,15 +19,15 @@ def __init__(
  self.filesystem = filesystem
  self.config = config
 
- def _save_dataframe(self, df, path, **kwargs) -> Optional[str]:
+ def _save_dataframe(self, df: pd.DataFrame, path: str, **kwargs) -> Optional[str]: # type: ignore
  errname = None
  try:
  self.filesystem.save_dataframe(df, path, **kwargs)
  except Exception as err:
  errname = f"Error during saving file {path}: {err}"
  return errname
 
- def validate_path_for_delete(self, columns_to_delete: List[str], path: str):
+ def validate_path_for_delete(self, columns_to_delete: List[str], path: str) -> None:
  df = self.filesystem.read_dataframe(path)
  for col in columns_to_delete:
  assert col in df.columns, f'Dataframe {path} dont have "{col}" column'
@@ -59,7 +59,7 @@ def delete_columns(
  )
  return [err for err in errors if err is not None]
 
- def validate_path_for_rename(self, column_map: Dict[str, str], path: str):
+ def validate_path_for_rename(self, column_map: Dict[str, str], path: str) -> None:
  df = self.filesystem.read_dataframe(path)
  for col_old, col_new in column_map.items():
  assert col_old in df.columns, f'Dataframe {path} dont have "{col_old}" column'
@@ -95,33 +95,32 @@ def rename_columns(
  def validate_path_for_update(
  self,
  key_column: str,
- df_new: List[dict],
+ df_new: List[Dict[str, Any]],
  path: str
- ):
+ ) -> None:
  df_new = pd.DataFrame(df_new)
  df_old = self.filesystem.read_dataframe(path)
  assert key_column in df_old.columns, f'Dataframe {path} dont have "{key_column}" column'
- assert set(df_old[key_column]) == set(df_new[key_column]), \
- f'Dataframe {path} has different values in "{key_column}"'
+ assert set(df_old[key_column]) == set(df_new[key_column]), f'Dataframe {path} has different values in "{key_column}"' # type: ignore
 
  duplicates = df_old[df_old[key_column].duplicated()][key_column].tolist()
  assert len(duplicates) == 0, f'Dataframe {path} has duplicates in "{key_column}" column: {duplicates}'
 
- duplicates = df_new[df_new[key_column].duplicated()][key_column].tolist()
+ duplicates = df_new[df_new[key_column].duplicated()][key_column].tolist() # type: ignore
  assert len(duplicates) == 0, f'New dataframe for {path} has duplicates in "{key_column}" column: {duplicates}'
 
  assert len(df_old) == len(df_new), f'Length of {path} dataframe is changed'
 
  def update_columns_for_path(
  self,
  key_column: str,
- df_new: List[dict],
+ df_new: List[Dict[str, Any]],
  path: str
  ) -> Optional[str]:
  df_new = pd.DataFrame(df_new)
  df_old = self.filesystem.read_dataframe(path)
 
- columns_to_add = [i for i in df_new.columns if i != key_column]
+ columns_to_add = [i for i in df_new.columns if i != key_column] # type: ignore [attr-defined]
  columns_intersection = set(df_old.columns).intersection(set(columns_to_add))
 
  if len(columns_intersection) > 0:
@@ -133,7 +132,7 @@ def update_columns_for_path(
  def update_columns(
  self,
  key_column: str,
- path2df: Dict[str, List[dict]],
+ path2df: Dict[str, List[Dict[str, Any]]],
  max_threads: int = 16,
  pbar: bool = True
  ) -> List[str]:

DPF/processors/processor.py

@@ -393,7 +393,7 @@ def save_to_shards(
  filesystem: Optional[FileSystem] = None,
  max_files_in_shard: int = 1000,
  datafiles_ext: str = "csv",
- archives_ext: Optional[str] = "tar",
+ archives_ext: str = "tar",
  filenaming: str = "counter",
  columns_to_save: Optional[List[str]] = None,
  rename_columns: Optional[Dict[str, str]] = None,
@@ -431,7 +431,7 @@ def save_to_shards(
  writer = ShardsWriter(
  filesystem,
  destination_dir,
- keys_mapping=rename_columns,
+ keys_to_rename=rename_columns,
  max_files_in_shard=max_files_in_shard,
  datafiles_ext=datafiles_ext,
  archives_ext=archives_ext,

DPF/processors/writers/filewriter.py

@@ -1,6 +1,6 @@
-import traceback
+from types import TracebackType
 from abc import abstractmethod
-from typing import Dict, Optional, Tuple
+from typing import Dict, Optional, Tuple, Union
 
 
 class ABSWriter:
@@ -19,8 +19,8 @@ def __enter__(self) -> "ABSWriter":
  @abstractmethod
  def __exit__(
  self,
- exception_type,
- exception_value: Optional[Exception],
- exception_traceback: traceback,
+ exception_type: Union[type[BaseException], None],
+ exception_value: Union[BaseException, None],
+ exception_traceback: Union[TracebackType, None],
  ) -> None:
  pass

DPF/processors/writers/sharded_files_writer.py

@@ -1,7 +1,8 @@
 import os
 import traceback
 import uuid
-from typing import Dict, List, Optional, Tuple
+from types import TracebackType
+from typing import Dict, List, Optional, Tuple, Any, Union
 
 import pandas as pd
 
@@ -34,9 +35,9 @@ def __init__(
  self.filenaming = filenaming
  assert self.filenaming in ["counter", "uuid"], "Invalid files naming"
 
- self.df_raw = []
+ self.df_raw: List[Dict[str, Any]] = []
  self.shard_index, self.last_file_index = self._init_writer_from_last_uploaded_file()
- self.last_path_to_dir = None
+ self.last_path_to_dir: str = None # type: ignore
 
  def save_sample(
  self,
@@ -66,20 +67,20 @@ def save_sample(
  self.df_raw.append(table_data)
  self._try_close_batch()
 
- def __enter__(self) -> "FileWriter": # noqa: F821
+ def __enter__(self) -> "ShardedFilesWriter": # noqa: F821
  return self
 
  def __exit__(
  self,
- exception_type: Optional[type],
- exception_value: Optional[Exception],
- exception_traceback: traceback,
+ exception_type: Union[type[BaseException], None],
+ exception_value: Union[BaseException, None],
+ exception_traceback: Union[TracebackType, None],
  ) -> None:
  if len(self.df_raw) != 0:
  self._flush(self._calculate_current_dirname())
  self.last_file_index = 0
 
- def _init_writer_from_last_uploaded_file(self) -> (int, int):
+ def _init_writer_from_last_uploaded_file(self) -> Tuple[int, int]:
  self.filesystem.mkdir(self.destination_dir)
  list_dirs = [
  int(os.path.basename(filename[: -len(self.datafiles_ext)]))
@@ -109,9 +110,12 @@ def _init_writer_from_last_uploaded_file(self) -> (int, int):
  def get_current_filename(self, extension: str) -> str:
  extension = extension.lstrip('.')
  if self.filenaming == "counter":
- return f"{self.last_file_index}.{extension}"
+ filename = f"{self.last_file_index}.{extension}"
  elif self.filenaming == "uuid":
- return f"{uuid.uuid4().hex}.{extension}"
+ filename = f"{uuid.uuid4().hex}.{extension}"
+ else:
+ raise ValueError(f"Invalid filenaming type: {self.filenaming}")
+ return filename
 
  def _calculate_current_dirname(self) -> str:
  return str(self.shard_index)

DPF/processors/writers/shards_writer.py

@@ -3,7 +3,8 @@
 import tarfile
 import traceback
 import uuid
-from typing import Dict, List, Optional, Tuple
+from types import TracebackType
+from typing import Dict, List, Optional, Tuple, Any, Union
 
 import pandas as pd
 
@@ -23,24 +24,24 @@ def __init__(
  self,
  filesystem: FileSystem,
  destination_dir: str,
- keys_mapping: Optional[Dict[str, str]] = None,
- max_files_in_shard: Optional[int] = 1000,
- datafiles_ext: Optional[str] = "csv",
- archives_ext: Optional[str] = "tar",
+ keys_to_rename: Optional[Dict[str, str]] = None,
+ max_files_in_shard: int = 1000,
+ datafiles_ext: str = "csv",
+ archives_ext: str = "tar",
  filenaming: str = "counter"
  ) -> None:
  self.filesystem = filesystem
  self.destination_dir = destination_dir
- self.keys_mapping = keys_mapping
+ self.keys_to_rename = keys_to_rename
  self.max_files_in_shard = max_files_in_shard
  self.datafiles_ext = "." + datafiles_ext.lstrip(".")
  self.archives_ext = "." + archives_ext.lstrip(".")
  self.filenaming = filenaming
  assert self.filenaming in ["counter", "uuid"], "Invalid files naming"
 
- self.df_raw = []
+ self.df_raw: List[Dict[str, Any]] = []
  self.tar_bytes = io.BytesIO()
- self.tar = None
+ self.tar: tarfile.TarFile = None # type: ignore
  self.shard_index, self.last_file_index = self._init_writer_from_last_uploaded_file()
 
  def save_sample(
@@ -61,8 +62,8 @@ def save_sample(
  img_tar_info, fp = self._prepare_image_for_tar_format(file_bytes, filename)
  self.tar.addfile(img_tar_info, fp)
 
- if self.keys_mapping:
- table_data = rename_dict_keys(table_data, self.keys_mapping)
+ if self.keys_to_rename:
+ table_data = rename_dict_keys(table_data, self.keys_to_rename)
 
  self.df_raw.append(table_data)
  self._try_close_batch()
@@ -76,20 +77,20 @@ def _prepare_image_for_tar_format(
  img_tar_info.size = len(fp.getvalue())
  return img_tar_info, fp
 
- def __enter__(self) -> "FileWriter": # noqa: F821
+ def __enter__(self) -> "ShardsWriter": # noqa: F821
  return self
 
  def __exit__(
  self,
- exception_type,
- exception_value: Optional[Exception],
- exception_traceback: traceback,
+ exception_type: Union[type[BaseException], None],
+ exception_value: Union[BaseException, None],
+ exception_traceback: Union[TracebackType, None],
  ) -> None:
  if len(self.df_raw) != 0:
  self._flush(self._calculate_current_tarname())
  self.last_file_index = 0
 
- def _init_writer_from_last_uploaded_file(self) -> (int, int):
+ def _init_writer_from_last_uploaded_file(self) -> Tuple[int, int]:
  self.filesystem.mkdir(self.destination_dir)
  list_csv = [
  int(os.path.basename(filename[: -len(self.datafiles_ext)]))
@@ -124,9 +125,12 @@ def _calculate_current_tarname(self) -> str:
  def get_current_filename(self, extension: str) -> str:
  extension = extension.lstrip('.')
  if self.filenaming == "counter":
- return f"{self.last_file_index}.{extension}"
+ filename = f"{self.last_file_index}.{extension}"
  elif self.filenaming == "uuid":
- return f"{uuid.uuid4().hex}.{extension}"
+ filename = f"{uuid.uuid4().hex}.{extension}"
+ else:
+ raise ValueError(f"Invalid filenaming type: {self.filenaming}")
+ return filename
 
  def _try_close_batch(self) -> None:
  old_tarname = self._calculate_current_tarname()
@@ -154,7 +158,7 @@ def _flush_and_upload_tar(self, filename: str) -> None:
  self.filesystem.save_file(
  self.tar_bytes, self.filesystem.join(self.destination_dir, filename), binary=True
  )
- self.tar = None
+ self.tar = None # type: ignore
  self.tar_bytes = io.BytesIO()
 
  def _flush(self, tarname: str) -> None:

DPF/processors/writers/utils.py

@@ -1,4 +1,7 @@
-def rename_dict_keys(d: dict, keys_mapping: dict) -> dict:
+from typing import Dict, Any
+
+
+def rename_dict_keys(d: Dict[Any, Any], keys_mapping: Dict[Any, Any]) -> Dict[Any, Any]:
  for k, v in keys_mapping.items():
  d[v] = d.pop(k)
  return d