From da771e0066c6b25946ec54471f5dad0f20f522a7 Mon Sep 17 00:00:00 2001
From: mark-oshea <mark.oshea96@rocketmail.com>
Date: Tue, 23 Jun 2026 16:36:44 -0700
Subject: [PATCH] Add Colorspace.BT2020 for BT.2020 YUV matrix conversion

---
 CHANGELOG.rst            |  1 +
 av/video/reformatter.pxd |  1 +
 av/video/reformatter.py  |  4 ++++
 av/video/reformatter.pyi |  2 ++
 tests/test_colorspace.py | 20 ++++++++++++++++++++
 5 files changed, 28 insertions(+)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 1e7a60595..b855a7fb5 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -41,6 +41,7 @@ Features:
 - Support ``yuv420p10le`` in ``VideoFrame.to_ndarray`` and ``VideoFrame.from_ndarray`` by :gh-user:`WyattBlue` (:issue:`1981`).
 - Add ``at`` parameter to ``Graph.push`` and ``Graph.vpush`` to push a frame to a single buffer source by index, for multi-input filters like ``overlay`` by :gh-user:`WyattBlue`.
 - ``find_best_pix_fmt_of_list`` now returns the loss as a ``PixFmtLoss`` ``enum.IntFlag`` instead of a plain ``int`` by :gh-user:`WyattBlue` (:issue:`2300`).
+- Add ``Colorspace.BT2020`` by :gh-user:`mark-oshea`.
 
 Fixes:
 
diff --git a/av/video/reformatter.pxd b/av/video/reformatter.pxd
index fe4b68db5..5315a7c10 100644
--- a/av/video/reformatter.pxd
+++ b/av/video/reformatter.pxd
@@ -33,6 +33,7 @@ cdef extern from "libswscale/swscale.h" nogil:
     cdef int SWS_CS_SMPTE170M
     cdef int SWS_CS_SMPTE240M
     cdef int SWS_CS_DEFAULT
+    cdef int SWS_CS_BT2020
 
     cdef SwsContext *sws_alloc_context()
     cdef void sws_free_context(SwsContext **ctx)
diff --git a/av/video/reformatter.py b/av/video/reformatter.py
index e95c871a3..eabcefe46 100644
--- a/av/video/reformatter.py
+++ b/av/video/reformatter.py
@@ -36,6 +36,7 @@ class Colorspace(IntEnum):
     SMPTE170M = SWS_CS_SMPTE170M
     SMPTE240M = SWS_CS_SMPTE240M
     DEFAULT = SWS_CS_DEFAULT
+    BT2020 = SWS_CS_BT2020
     # Lowercase for b/c.
     itu709 = SWS_CS_ITU709
     fcc = SWS_CS_FCC
@@ -44,6 +45,7 @@ class Colorspace(IntEnum):
     smpte170m = SWS_CS_SMPTE170M
     smpte240m = SWS_CS_SMPTE240M
     default = SWS_CS_DEFAULT
+    bt2020 = SWS_CS_BT2020
 
 
 class ColorRange(IntEnum):
@@ -146,6 +148,8 @@ def _set_frame_colorspace(
         frame.colorspace = lib.AVCOL_SPC_SMPTE170M
     elif colorspace == SWS_CS_SMPTE240M:
         frame.colorspace = lib.AVCOL_SPC_SMPTE240M
+    elif colorspace == SWS_CS_BT2020:
+        frame.colorspace = lib.AVCOL_SPC_BT2020_NCL
 
 
 @cython.final
diff --git a/av/video/reformatter.pyi b/av/video/reformatter.pyi
index 30508a32b..480860d0b 100644
--- a/av/video/reformatter.pyi
+++ b/av/video/reformatter.pyi
@@ -31,6 +31,7 @@ class Colorspace(IntEnum):
     SMPTE170M = cast(int, ...)
     SMPTE240M = cast(int, ...)
     DEFAULT = cast(int, ...)
+    BT2020 = cast(int, ...)
     itu709 = cast(int, ...)
     fcc = cast(int, ...)
     itu601 = cast(int, ...)
@@ -38,6 +39,7 @@ class Colorspace(IntEnum):
     smpte170m = cast(int, ...)
     smpte240m = cast(int, ...)
     default = cast(int, ...)
+    bt2020 = cast(int, ...)
 
 class ColorRange(IntEnum):
     UNSPECIFIED = 0
diff --git a/tests/test_colorspace.py b/tests/test_colorspace.py
index a5352fc7d..6afd42cf0 100644
--- a/tests/test_colorspace.py
+++ b/tests/test_colorspace.py
@@ -1,3 +1,5 @@
+import pytest
+
 import av
 from av.video.reformatter import (
     ColorPrimaries,
@@ -101,3 +103,21 @@ def test_reformat_preserves_color_primaries() -> None:
     frame.color_primaries = ColorPrimaries.BT709
     rgb = frame.reformat(format="rgb24")
     assert rgb.color_primaries == ColorPrimaries.BT709
+
+
+@pytest.mark.parametrize(
+    ("colorspace", "expected"),
+    [
+        (Colorspace.ITU709, Colorspace.ITU709),
+        (Colorspace.FCC, Colorspace.FCC),
+        (Colorspace.ITU601, 6),  # AVCOL_SPC_SMPTE170M
+        (Colorspace.SMPTE240M, Colorspace.SMPTE240M),
+        (Colorspace.BT2020, Colorspace.BT2020),
+    ],
+)
+def test_reformat_dst_colorspace_metadata(
+    colorspace: Colorspace, expected: Colorspace | int
+) -> None:
+    frame = av.VideoFrame(width=64, height=64, format="yuv420p")
+    rgb = frame.reformat(format="rgb24", dst_colorspace=colorspace)
+    assert rgb.colorspace == expected