refactor: ♻️ change FILE_FORMAT to enum.Enum

IRFM · Mar 11, 2024 · da5b5cc · da5b5cc
1 parent 9719911
commit da5b5cc
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 18 deletions.
diff --git a/src/python/librir/video_io/IRMovie.py b/src/python/librir/video_io/IRMovie.py
@@ -12,7 +12,6 @@
 from librir.tools.FileAttributes import FileAttributes
 from librir.tools._thermavip import init_thermavip, unbind_thermavip_shared_mem
 from librir.video_io.rir_video_io import (
-    FILE_FORMAT_H264,
     enable_motion_correction,
     load_motion_correction_file,
     motion_correction_enabled,
@@ -21,6 +20,7 @@
 
 from .IRSaver import IRSaver
 from .rir_video_io import (
+    FileFormat,
     calibrate_image,
     calibration_files,
     close_camera,
@@ -536,7 +536,7 @@ def _build_outfile(self) -> str:
         parent, basename = os.path.split(f)
         stem = basename.replace(suffix, "")
 
-        if self.video_file_format != FILE_FORMAT_H264:
+        if self.video_file_format != FileFormat.H264:
             return os.path.abspath((os.path.join(parent, (stem + ".h264"))))
 
         return self.filename

diff --git a/src/python/librir/video_io/__init__.py b/src/python/librir/video_io/__init__.py
@@ -6,11 +6,6 @@
 
 # import useful functions from rir_video_io
 from .rir_video_io import (
-    FILE_FORMAT_H264,
-    FILE_FORMAT_PCR,
-    FILE_FORMAT_PCR_ENCAPSULATED,
-    FILE_FORMAT_WEST,
-    FILE_FORMAT_ZSTD_COMPRESSED,
     correct_PCR_file,
     video_file_format,
     get_filename,
@@ -23,10 +18,6 @@
     "IRMovie",
     "is_ir_file_corrupted",
     "video_file_format",
-    "FILE_FORMAT_PCR",
-    "FILE_FORMAT_WEST",
-    "FILE_FORMAT_PCR_ENCAPSULATED",
-    "FILE_FORMAT_ZSTD_COMPRESSED",
-    "FILE_FORMAT_H264",
+    "FileFormat",
     "get_filename",
 ]
diff --git a/src/python/librir/video_io/rir_video_io.py b/src/python/librir/video_io/rir_video_io.py
@@ -2,19 +2,20 @@
 import logging
 
 import numpy as np
-
+import enum
 from ..low_level.misc import _video_io, toArray, toBytes, toString
 
 logging.basicConfig()
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
 
-FILE_FORMAT_PCR = 1
-FILE_FORMAT_WEST = 2
-FILE_FORMAT_PCR_ENCAPSULATED = 3
-FILE_FORMAT_ZSTD_COMPRESSED = 4
-FILE_FORMAT_H264 = 5
+class FileFormat(enum.Enum):
+    PCR = 1
+    WEST = 2
+    PCR_ENCAPSULATED = 3
+    ZSTD_COMPRESSED = 4
+    H264 = 5
 
 
 def open_camera_file(filename):