allow enum and literal to f3kdb constructor

Jaded-Encoding-Thaumaturgy · Jul 24, 2021 · 5d6c935 · 5d6c935
1 parent 99dfc1b
commit 5d6c935
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/debandshit/debanders.py b/debandshit/debanders.py
@@ -9,7 +9,7 @@
 import vapoursynth as vs
 from vsutil import depth
 
-from .f3kdb import SAMPLEMODE, F3kdb
+from .f3kdb import SAMPLEMODE, F3kdb, SampleMode
 from .placebo import Placebo
 
 core = vs.core
@@ -19,7 +19,7 @@
 
 def dumb3kdb(clip: vs.VideoNode, radius: int = 16,
              threshold: Union[int, List[int]] = 30, grain: Union[int, List[int]] = 0,
-             sample_mode: SAMPLEMODE = 2, use_neo: bool = False, **kwargs: Any) -> vs.VideoNode:
+             sample_mode: Union[SAMPLEMODE, SampleMode] = 2, use_neo: bool = False, **kwargs: Any) -> vs.VideoNode:
     """Small convenience function for calling F3kdb().deband()."""
     return F3kdb(radius, threshold, grain, sample_mode, use_neo, **kwargs).deband(clip)
 

diff --git a/debandshit/f3kdb.py b/debandshit/f3kdb.py
@@ -1,16 +1,24 @@
+from enum import IntEnum
 from typing import Any, Dict, List, Literal, Union
 
 import vapoursynth as vs
 
 core = vs.core
 
 
-__all__ = ['SAMPLEMODE', 'F3kdb']
+__all__ = ['SAMPLEMODE', 'F3kdb', 'SampleMode']
 
 
 SAMPLEMODE = Literal[1, 2, 3, 4]
 
 
+class SampleMode(IntEnum):
+    COLUMN = 1
+    SQUARE = 2
+    ROW = 3
+    COL_ROW_MEAN = 4
+
+
 class F3kdb:
     """f3kdb object."""
     radius: int
@@ -19,7 +27,7 @@ class F3kdb:
     thcr: int
     gry: int
     grc: int
-    sample_mode: SAMPLEMODE
+    sample_mode: int
     use_neo: bool
     f3kdb_args: Dict[str, Any]
 
@@ -28,7 +36,7 @@ class F3kdb:
     def __init__(self,
                  radius: int = 16,
                  threshold: Union[int, List[int]] = 30, grain: Union[int, List[int]] = 0,
-                 sample_mode: SAMPLEMODE = 2, use_neo: bool = False, **kwargs: Any) -> None:
+                 sample_mode: Union[SAMPLEMODE, SampleMode] = 2, use_neo: bool = False, **kwargs: Any) -> None:
         """
         Handle debanding operations onto a clip using a set of configured parameters.
 
@@ -41,14 +49,14 @@ def __init__(self,
         :param threshold:       Banding detection threshold(s) for planes
         :param grain:           Specifies amount of grains added in the last debanding stage
         :param sample_mode:     Valid modes are:
-                                * SampleMode.COLUMN: Take 2 pixels as reference pixel
+                                * SampleMode.COLUMN or 1: Take 2 pixels as reference pixel
                                   Reference pixels are in the same column of current pixel
-                                * SampleMode.SQUARE: Take 4 pixels as reference pixel
+                                * SampleMode.SQUARE or 2: Take 4 pixels as reference pixel
                                   Reference pixels are in the square around current pixel
-                                * SampleMode.ROW: Take 2 pixels as reference pixel
+                                * SampleMode.ROW or 3: Take 2 pixels as reference pixel
                                   Reference pixels are in the same row of current pixel
                                   Only `neo_f3kdb.Deband` supports it
-                                * SampleMode.COL_ROW_MEAN: Arithmetic mean of 1 and 3
+                                * SampleMode.COL_ROW_MEAN or 4: Arithmetic mean of 1 and 3
                                   Reference points are randomly picked within the range
                                   Only `neo_f3kdb.Deband` supports it
         :param use_neo:         Use `neo_f3kdb.Deband`