1 files changed, 70 insertions, 16 deletions
diff --git a/py/tests/gamma.py b/py/tests/gamma.py
index 5969b82..03499bc 100755
--- a/py/tests/gamma.py
+++ b/py/tests/gamma.py
@@ -1,14 +1,16 @@
 #!/usr/bin/python3
 
 import pykms
+import argparse
 
-# This hack makes drm initialize the fbcon, setting up the default connector
-card = pykms.Card()
-card = 0
+parser = argparse.ArgumentParser()
+parser.add_argument("-c", "--connector", default="")
+parser.add_argument("-l", "--legacy", action="store_true", default=False)
+args = parser.parse_args()
 
 card = pykms.Card()
 res = pykms.ResourceManager(card)
-conn = res.reserve_connector()
+conn = res.reserve_connector(args.connector)
 crtc = res.reserve_crtc(conn)
 mode = conn.get_default_mode()
 
@@ -17,24 +19,76 @@ pykms.draw_test_pattern(fb);
 
 crtc.set_mode(conn, fb, mode)
 
-len=256
-arr = bytearray(len*2*4)
-view = memoryview(arr).cast("H")
+use_legacy = args.legacy
+
+if not use_legacy:
+	prop = crtc.get_prop("GAMMA_LUT")
+
+	if not prop:
+		prop = crtc.get_prop("DEGAMMA_LUT")
+
+		if not prop:
+			print("No gamma property found")
+			exit(-1)
+		else:
+			print("Using DEGAMMA_LUT for gamma")
+
+
+def legacy_gamma_set():
+	len = crtc.legacy_gamma_size()
+
+	table = []
+
+	for i in range(len):
+	    g = round(65535 * pow(i / float(len), 1 / 2.2))
+	    table.append((g, g, g))
+
+	crtc.legacy_gamma_set(table)
+
+def legacy_gamma_clear():
+	len = crtc.legacy_gamma_size()
+
+	table = []
+
+	for i in range(len):
+	    g = round(65535 * (i / float(len)))
+	    table.append((g, g, g))
+
+	crtc.legacy_gamma_set(table)
+
+def gamma_set():
+	len=256
+	arr = bytearray(len*2*4)
+	view = memoryview(arr).cast("H")
+
+	for i in range(len):
+	    g = round(65535 * pow(i / float(len), 1 / 2.2))
+
+	    view[i * 4 + 0] = g
+	    view[i * 4 + 1] = g
+	    view[i * 4 + 2] = g
+	    view[i * 4 + 3] = 0
+
+	gamma = pykms.Blob(card, arr);
+
+	crtc.set_prop(prop, gamma.id)
 
-for i in range(len):
-    g = round(65535 * pow(i / float(len), 1 / 2.2))
 
-    view[i * 4 + 0] = g
-    view[i * 4 + 1] = g
-    view[i * 4 + 2] = g
-    view[i * 4 + 3] = 0
+def gamma_clear():
+	crtc.set_prop(prop, 0)
 
-gamma = pykms.Blob(card, arr);
+input("press enter to apply gamma\n")
 
-crtc.set_prop("GAMMA_LUT", gamma.id)
+if use_legacy:
+	legacy_gamma_set()
+else:
+	gamma_set()
 
 input("press enter to remove gamma\n")
 
-crtc.set_prop("GAMMA_LUT", 0)
+if use_legacy:
+	legacy_gamma_clear()
+else:
+	gamma_clear()
 
 input("press enter to exit\n")