Some fixes, re-wording, better testing. (#28)

* Check whether get_y is passed a yuv-clip * Update vsutil.py * Fix some docstring formatting, grammar, capitalization * Re-organize imports according to PEP 8 Since I expect most users to be importing `*` from this module, we should be using absolute imports for the few standard library functions we need. The alternative to this would be putting the os/mimetypes imports within the `is_image` function, as this would prevent the `from vsutil import *` user from having these modules' namespaces reserved, but this is not the proper location for imports according to PEP 8. * Fix missing f-strings in tests, remove default values in function calls Additionally, change `_format` to `format` in BlankClip params as this is how it is documented. * Cleanup get_subsampling, check for YUV, add RaisesRegex tests * Re-word error message in get_y, add RaisesRegex test * Formatting get_w No need to int(round(width)) as round() with no second argument will always return an int. Also no need to re-round width after rounding to nearest even int, as interger multiplication will also result in a int. * Update LICENSE, fix VapourSynth capitalization in README.rst * Add __all__ list for the `from vsutil import *` wild card user This allows cleaner looking scripts, and might as well be the recommended way of import from vsutil _in scripts_ not other modules. With `core` and `vs` being defined in __all__, a VS script will not need both of the standard VS lines if using the wild card import. * Change remaining double quotes to single quotes * get_subsampling: return None instead of an empty str for non-YUV formats * Fix join `planes` list to grab all three clips' first plane It is documented being a list [0, 0, 0]: http://www.vapoursynth.com/doc/functions/shuffleplanes.html * Use Python 3.8's positional-only parameters syntax and Literals https://www.python.org/dev/peps/pep-0570/ for more information on positional-only parameters. https://www.python.org/dev/peps/pep-0586/ for more information on typing.Literal. * Fix weird Literal implementation spec warning with vs.RGB etc. According to https://www.python.org/dev/peps/pep-0586/ , Literals can be parameterized with enums, but `vs.RGB` isn't *actually* an enum in Python's eyes (even though it's technically exposed as one). `vs.ColorFamily.RGB` is the correct way of handling this enum and both refer to the same object (i.e. `vs.ColorFamily.RGB is vs.RGB`). * make it possible to pass height as kwarg to get_w Co-authored-by: Roland Netzsch <[email protected]> Co-authored-by: kageru <[email protected]>
Irrational-Encoding-Wizardry · May 26, 2020 · dc3ed6a · dc3ed6a
1 parent c0206e2
commit dc3ed6a
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 73 deletions.
diff --git a/LICENSE b/LICENSE
@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2019 Irrational Encoding Wizardry
+Copyright (c) 2020 Irrational Encoding Wizardry
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

diff --git a/README.rst b/README.rst
@@ -6,19 +6,19 @@ vsutil
 .. |circleci| image:: https://img.shields.io/circleci/project/github/Irrational-Encoding-Wizardry/vsutil.svg
   :target: https://circleci.com/gh/Irrational-Encoding-Wizardry/vsutil
   :alt: CircleCI (all branches)
-  
+
 .. |codecov| image:: https://img.shields.io/codecov/c/gh/Irrational-Encoding-Wizardry/vsutil.svg
   :target: https://codecov.io/gh/Irrational-Encoding-Wizardry/vsutil
   :alt: Codecov
-  
+
 .. |discord| image:: https://img.shields.io/discord/221919789017202688.svg
   :target: https://discord.gg/ZB7ZXbN
   :alt: Discord
 
-A collection of general-purpose Vapoursynth functions to be reused in modules and scripts.
+A collection of general-purpose VapourSynth functions to be reused in modules and scripts.
 
-The goal for vsutil is to allow authors of various "func" scripts to make use of premade helper functions instead of having to write their own. 
+The goal for vsutil is to allow authors of various "func" scripts to make use of premade helper functions instead of having to write their own.
 
-There are various benefits to this. For starters, only one script will require updating if anything is changed in Vapoursynth instead of every single func. This also helps unmaintained scripts from breaking and never being fixed. Additionally this will also be a good resource for new authors and makes functions easier to write and read.
+There are various benefits to this. For starters, only one script will require updating if anything is changed in VapourSynth instead of every single func. This also helps unmaintained scripts from breaking and never being fixed. Additionally this will also be a good resource for new authors and makes functions easier to write and read.
 
-As this is a community-driven project, contributions are heavily encouraged. vsutil will be continually updated to ensure it is up-to-date with changes to Vapoursynth and to include various pull requests sent in by contributors.
+As this is a community-driven project, contributions are heavily encouraged. vsutil will be continually updated to ensure it is up-to-date with changes to VapourSynth and to include various pull requests sent in by contributors.
diff --git a/tests/test_vsutil.py b/tests/test_vsutil.py
@@ -4,44 +4,45 @@
 
 
 class VsUtilTests(unittest.TestCase):
-    YUV420P8_CLIP = vs.core.std.BlankClip(_format=vs.YUV420P8, width=160, height=120, color=[0, 128, 128], length=100)
-    YUV420P10_CLIP = vs.core.std.BlankClip(_format=vs.YUV420P10, width=160, height=120, color=[0, 128, 128], length=100)
-    YUV444P8_CLIP = vs.core.std.BlankClip(_format=vs.YUV444P8, width=160, height=120, color=[0, 128, 128], length=100)
-    YUV422P8_CLIP = vs.core.std.BlankClip(_format=vs.YUV422P8, width=160, height=120, color=[0, 128, 128], length=100)
-    YUV410P8_CLIP = vs.core.std.BlankClip(_format=vs.YUV410P8, width=160, height=120, color=[0, 128, 128], length=100)
-    YUV411P8_CLIP = vs.core.std.BlankClip(_format=vs.YUV411P8, width=160, height=120, color=[0, 128, 128], length=100)
-    YUV440P8_CLIP = vs.core.std.BlankClip(_format=vs.YUV440P8, width=160, height=120, color=[0, 128, 128], length=100)
-
-    SMALLER_SAMPLE_CLIP = vs.core.std.BlankClip(_format=vs.YUV420P8, width=10, height=10)
-
-    BLACK_SAMPLE_CLIP = vs.core.std.BlankClip(_format=vs.YUV420P8, width=160, height=120, color=[0, 128, 128],
+    YUV420P8_CLIP = vs.core.std.BlankClip(format=vs.YUV420P8, width=160, height=120, color=[0, 128, 128], length=100)
+    YUV420P10_CLIP = vs.core.std.BlankClip(format=vs.YUV420P10, width=160, height=120, color=[0, 128, 128], length=100)
+    YUV444P8_CLIP = vs.core.std.BlankClip(format=vs.YUV444P8, width=160, height=120, color=[0, 128, 128], length=100)
+    YUV422P8_CLIP = vs.core.std.BlankClip(format=vs.YUV422P8, width=160, height=120, color=[0, 128, 128], length=100)
+    YUV410P8_CLIP = vs.core.std.BlankClip(format=vs.YUV410P8, width=160, height=120, color=[0, 128, 128], length=100)
+    YUV411P8_CLIP = vs.core.std.BlankClip(format=vs.YUV411P8, width=160, height=120, color=[0, 128, 128], length=100)
+    YUV440P8_CLIP = vs.core.std.BlankClip(format=vs.YUV440P8, width=160, height=120, color=[0, 128, 128], length=100)
+    RGB24_CLIP = vs.core.std.BlankClip(format=vs.RGB24)
+
+    SMALLER_SAMPLE_CLIP = vs.core.std.BlankClip(format=vs.YUV420P8, width=10, height=10)
+
+    BLACK_SAMPLE_CLIP = vs.core.std.BlankClip(format=vs.YUV420P8, width=160, height=120, color=[0, 128, 128],
                                               length=100)
-    WHITE_SAMPLE_CLIP = vs.core.std.BlankClip(_format=vs.YUV420P8, width=160, height=120, color=[255, 128, 128],
+    WHITE_SAMPLE_CLIP = vs.core.std.BlankClip(format=vs.YUV420P8, width=160, height=120, color=[255, 128, 128],
                                               length=100)
 
     def assert_same_dimensions(self, clip_a: vs.VideoNode, clip_b: vs.VideoNode):
         """
         Assert that two clips have the same width and height.
         """
-        self.assertEqual(clip_a.height, clip_b.height, 'Same height expected, was {clip_a.height} and {clip_b.height}')
-        self.assertEqual(clip_a.width, clip_b.width, 'Same width expected, was {clip_a.width} and {clip_b.width}')
+        self.assertEqual(clip_a.height, clip_b.height, f'Same height expected, was {clip_a.height} and {clip_b.height}.')
+        self.assertEqual(clip_a.width, clip_b.width, f'Same width expected, was {clip_a.width} and {clip_b.width}.')
 
     def assert_same_format(self, clip_a: vs.VideoNode, clip_b: vs.VideoNode):
         """
         Assert that two clips have the same format (but not necessarily size).
         """
-        self.assertEqual(clip_a.format.id, clip_b.format.id, 'Same format expected')
+        self.assertEqual(clip_a.format.id, clip_b.format.id, 'Same format expected.')
 
     def assert_same_bitdepth(self, clip_a: vs.VideoNode, clip_b: vs.VideoNode):
         """
         Assert that two clips have the same number of bits per sample.
         """
         self.assertEqual(clip_a.format.bits_per_sample, clip_b.format.bits_per_sample,
-                         'Same depth expected, was {clip_a.format.bits_per_sample} and {clip_b.format.bits_per_sample}')
+                         f'Same depth expected, was {clip_a.format.bits_per_sample} and {clip_b.format.bits_per_sample}.')
 
     def assert_same_length(self, clip_a: vs.VideoNode, clip_b: vs.VideoNode):
         self.assertEqual(len(clip_a), len(clip_b),
-                         'Same number of frames expected, was {len(clip_a)} and {len(clip_b)}.')
+                         f'Same number of frames expected, was {len(clip_a)} and {len(clip_b)}.')
 
     def assert_same_metadata(self, clip_a: vs.VideoNode, clip_b: vs.VideoNode):
         """
@@ -66,8 +67,9 @@ def test_subsampling(self):
         self.assertEqual('422', vsutil.get_subsampling(self.YUV422P8_CLIP))
         self.assertEqual('411', vsutil.get_subsampling(self.YUV411P8_CLIP))
         self.assertEqual('410', vsutil.get_subsampling(self.YUV410P8_CLIP))
+        self.assertEqual(None, vsutil.get_subsampling(self.RGB24_CLIP))
         # let’s create a custom format with higher subsampling than any of the legal ones to test that branch as well:
-        with self.assertRaises(ValueError):
+        with self.assertRaisesRegex(ValueError, 'Unknown subsampling.'):
             vsutil.get_subsampling(
                 vs.core.std.BlankClip(_format=self.YUV444P8_CLIP.format.replace(subsampling_w=4))
             )
@@ -120,6 +122,9 @@ def test_get_y(self):
         self.assert_same_dimensions(self.BLACK_SAMPLE_CLIP, y)
         self.assert_same_bitdepth(self.BLACK_SAMPLE_CLIP, y)
 
+        with self.assertRaisesRegex(ValueError, 'The clip must have a luma plane.'):
+            vsutil.get_y(self.RGB24_CLIP)
+
     def test_split_join(self):
         planes = vsutil.split(self.BLACK_SAMPLE_CLIP)
         self.assertEqual(len(planes), 3)
@@ -133,7 +138,7 @@ def test_frame2clip(self):
         try:
             vs.core.add_cache = False
             black_frame = self.BLACK_SAMPLE_CLIP.get_frame(0)
-            black_clip = vsutil.frame2clip(black_frame, enforce_cache=True)
+            black_clip = vsutil.frame2clip(black_frame)
             self.assert_same_frame(self.BLACK_SAMPLE_CLIP, black_clip)
         # reset state of the core for further tests
         finally:
@@ -145,10 +150,10 @@ def test_is_image(self):
         self.assertEqual(vsutil.is_image('something.m2ts'), False)
 
     def test_get_w(self):
-        self.assertEqual(vsutil.get_w(480, only_even=True), 854)
+        self.assertEqual(vsutil.get_w(480), 854)
         self.assertEqual(vsutil.get_w(480, only_even=False), 853)
         self.assertEqual(vsutil.get_w(1080, 4 / 3), 1440)
-        self.assertEqual(vsutil.get_w(1080, 16 / 9), 1920)
+        self.assertEqual(vsutil.get_w(1080), 1920)
 
     def test_iterate(self):
         def double_number(x: int) -> int:

diff --git a/vsutil.py b/vsutil.py
@@ -1,58 +1,64 @@
 """
-VSUtil. A collection of general-purpose Vapoursynth functions to be reused in modules and scripts.
+VSUtil. A collection of general-purpose VapourSynth functions to be reused in modules and scripts.
 """
+__all__ = ['core', 'fallback', 'frame2clip', 'get_depth', 'get_plane_size', 'get_subsampling', 'get_w', 'get_y',
+           'insert_clip', 'is_image', 'iterate', 'join', 'plane', 'split', 'vs']
+
 from functools import reduce
-from typing import Callable, TypeVar, Union, List, Tuple, Optional
+from mimetypes import types_map
+from os import path
+from typing import Callable, List, Literal, Optional, Tuple, TypeVar, Union
+
 import vapoursynth as vs
-import mimetypes
-import os
 
 core = vs.core
-T = TypeVar("T")
+T = TypeVar('T')
 
 
-def get_subsampling(clip: vs.VideoNode) -> str:
+def get_subsampling(clip: vs.VideoNode, /) -> Union[None, str]:
     """
-    Returns the subsampling of a clip in human-readable format.
+    Returns the subsampling of a VideoNode in human-readable format.
+    Returns None for formats without subsampling.
     """
+    if clip.format.color_family not in (vs.YUV, vs.YCOCG):
+        return None
     if clip.format.subsampling_w == 1 and clip.format.subsampling_h == 1:
-        css = '420'
+        return '420'
     elif clip.format.subsampling_w == 1 and clip.format.subsampling_h == 0:
-        css = '422'
+        return '422'
     elif clip.format.subsampling_w == 0 and clip.format.subsampling_h == 0:
-        css = '444'
+        return '444'
     elif clip.format.subsampling_w == 2 and clip.format.subsampling_h == 2:
-        css = '410'
+        return '410'
     elif clip.format.subsampling_w == 2 and clip.format.subsampling_h == 0:
-        css = '411'
+        return '411'
     elif clip.format.subsampling_w == 0 and clip.format.subsampling_h == 1:
-        css = '440'
+        return '440'
     else:
-        raise ValueError('Unknown subsampling')
-    return css
+        raise ValueError('Unknown subsampling.')
 
 
-def get_depth(clip: vs.VideoNode) -> int:
+def get_depth(clip: vs.VideoNode, /) -> int:
     """
-    Returns the bitdepth of a clip as an integer.
+    Returns the bit depth of a VideoNode as an integer.
     """
     return clip.format.bits_per_sample
 
 
-def get_plane_size(frame: Union[vs.VideoFrame, vs.VideoNode], planeno: int) -> Tuple[int, int]:
+def get_plane_size(frame: Union[vs.VideoFrame, vs.VideoNode], /, planeno: int) -> Tuple[int, int]:
     """
-    Calculates the size of the plane
-    
-    :param frame:    The frame
-    :param planeno:  The plane
+    Calculates the dimensions (w, h) of the desired plane.
+
+    :param frame:    Can be a clip or frame.
+    :param planeno:  The desired plane's index.
     :return: (width, height)
     """
-    # Add additional checks on Video Nodes as their size and format can be variable.
+    # Add additional checks on VideoNodes as their size and format can be variable.
     if isinstance(frame, vs.VideoNode):
         if frame.width == 0:
-            raise ValueError("Cannot calculate plane size of variable size clip. Pass a frame instead.")
+            raise ValueError('Cannot calculate plane size of variable size clip. Pass a frame instead.')
         if frame.format is None:
-            raise ValueError("Cannot calculate plane size of variable format clip. Pass a frame instead.")
+            raise ValueError('Cannot calculate plane size of variable format clip. Pass a frame instead.')
 
     width, height = frame.width, frame.height
     if planeno != 0:
@@ -63,12 +69,12 @@ def get_plane_size(frame: Union[vs.VideoFrame, vs.VideoNode], planeno: int) -> T
 
 def iterate(base: T, function: Callable[[T], T], count: int) -> T:
     """
-    Utility function that executes a given function for a given number of times.
+    Utility function that executes a given function a given number of times.
     """
     return reduce(lambda v, _: function(v), range(count), base)
 
 
-def insert_clip(clip: vs.VideoNode, insert: vs.VideoNode, start_frame: int) -> vs.VideoNode:
+def insert_clip(clip: vs.VideoNode, /, insert: vs.VideoNode, start_frame: int) -> vs.VideoNode:
     """
     Convenience method to insert a shorter clip into a longer one.
     The inserted clip cannot go beyond the last frame of the source clip or an exception is raised.
@@ -78,7 +84,7 @@ def insert_clip(clip: vs.VideoNode, insert: vs.VideoNode, start_frame: int) -> v
     pre = clip[:start_frame]
     frame_after_insert = start_frame + insert.num_frames
     if frame_after_insert > clip.num_frames:
-        raise ValueError('Inserted clip is too long')
+        raise ValueError('Inserted clip is too long.')
     if frame_after_insert == clip.num_frames:
         return pre + insert
     post = clip[start_frame + insert.num_frames:]
@@ -92,45 +98,49 @@ def fallback(value: Optional[T], fallback_value: T) -> T:
     return fallback_value if value is None else value
 
 
-def plane(clip: vs.VideoNode, planeno: int) -> vs.VideoNode:
+def plane(clip: vs.VideoNode, planeno: int, /) -> vs.VideoNode:
     """
     Extract the plane with the given index from the clip.
-    
+
     :param clip:     The clip to extract the plane from.
-    :param planeno:  The planeno that specifies which plane to extract.
+    :param planeno:  The index that specifies which plane to extract.
     :return: A grayscale clip that only contains the given plane.
     """
     return core.std.ShufflePlanes(clip, planeno, vs.GRAY)
 
 
-def get_y(clip: vs.VideoNode) -> vs.VideoNode:
+def get_y(clip: vs.VideoNode, /) -> vs.VideoNode:
     """
     Helper to get the luma of a VideoNode.
     """
+    if clip.format is None or clip.format.color_family not in (vs.YUV, vs.YCOCG):
+        raise ValueError('The clip must have a luma plane.')
+
     return plane(clip, 0)
 
 
-def split(clip: vs.VideoNode) -> List[vs.VideoNode]:
+def split(clip: vs.VideoNode, /) -> List[vs.VideoNode]:
     """
     Returns a list of planes for the given input clip.
     """
     return [plane(clip, x) for x in range(clip.format.num_planes)]
 
 
-def join(planes: List[vs.VideoNode], family=vs.YUV) -> vs.VideoNode:
+def join(planes: List[vs.VideoNode],
+         family: Literal[vs.ColorFamily.RGB, vs.ColorFamily.YUV, vs.ColorFamily.YCOCG] = vs.YUV) -> vs.VideoNode:
     """
-    Joins the supplied list of planes into a YUV video node.
+    Joins the supplied list of planes into a three-plane VideoNode (defaults to YUV).
     """
-    return core.std.ShufflePlanes(clips=planes, planes=[0], colorfamily=family)
+    return core.std.ShufflePlanes(clips=planes, planes=[0, 0, 0], colorfamily=family)
 
 
-def frame2clip(frame: vs.VideoFrame, *, enforce_cache=True) -> vs.VideoNode:
+def frame2clip(frame: vs.VideoFrame, /, *, enforce_cache=True) -> vs.VideoNode:
     """
-    Converts a vapoursynth frame to a clip.
-    
+    Converts a VapourSynth frame to a clip.
+
     :param frame:         The frame to wrap.
     :param enforce_cache: Always add a cache. (Even if the vapoursynth module has this feature disabled)
-    :returns:     A one-frame VideoNode that yields the frame passed to the function.
+    :return: A one-frame VideoNode that yields the frame passed to the function.
     """
     bc = core.std.BlankClip(
         width=frame.width,
@@ -150,19 +160,19 @@ def frame2clip(frame: vs.VideoFrame, *, enforce_cache=True) -> vs.VideoNode:
     return result
 
 
-def get_w(height: int, aspect_ratio: float = 16 / 9, only_even: bool = True) -> int:
+def get_w(height: int, aspect_ratio: float = 16 / 9, *, only_even: bool = True) -> int:
     """
     Calculates the width for a clip with the given height and aspect ratio.
     only_even is True by default because it imitates the math behind most standard resolutions (e.g. 854x480).
     """
     width = height * aspect_ratio
     if only_even:
-        width = round(width / 2) * 2
-    return int(round(width))
+        return round(width / 2) * 2
+    return round(width)
 
 
-def is_image(filename: str) -> bool:
+def is_image(filename: str, /) -> bool:
     """
     Returns true if a filename refers to an image.
     """
-    return mimetypes.types_map.get(os.path.splitext(filename)[-1], "").startswith("image/")
+    return types_map.get(path.splitext(filename)[-1], '').startswith('image/')