CREATE EXTENSION jsonb_plpython2u CASCADE;
NOTICE:  installing required extension "plpython2u"
-- test jsonb -> python dict
CREATE FUNCTION test1(val jsonb) RETURNS int
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
assert isinstance(val, dict)
assert(val == {'a': 1, 'c': 'NULL'})
return len(val)
$$;
SELECT test1('{"a": 1, "c": "NULL"}'::jsonb);
 test1 
-------
     2
(1 row)

-- test jsonb -> python dict
-- complex dict with dicts as value
CREATE FUNCTION test1complex(val jsonb) RETURNS int
LANGUAGE plpython2u
TRANSFORM FOR TYPE jsonb
AS $$
assert isinstance(val, dict)
assert(val == {"d": {"d": 1}})
return len(val)
$$;
SELECT test1complex('{"d": {"d": 1}}'::jsonb);
 test1complex 
--------------
            1
(1 row)

-- test jsonb[] -> python dict
-- dict with array as value
CREATE FUNCTION test1arr(val jsonb) RETURNS int
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
assert isinstance(val, dict)
assert(val == {"d": [12, 1]})
return len(val)
$$;
SELECT test1arr('{"d":[12, 1]}'::jsonb);
 test1arr 
----------
        1
(1 row)

-- test jsonb[] -> python list
-- simple list
CREATE FUNCTION test2arr(val jsonb) RETURNS int
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
assert isinstance(val, list)
assert(val == [12, 1])
return len(val)
$$;
SELECT test2arr('[12, 1]'::jsonb);
 test2arr 
----------
        2
(1 row)

-- test jsonb[] -> python list
-- array of dicts
CREATE FUNCTION test3arr(val jsonb) RETURNS int
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
assert isinstance(val, list)
assert(val == [{"a": 1,"b": 2}, {"c": 3,"d": 4}])
return len(val)
$$;
SELECT test3arr('[{"a": 1, "b": 2}, {"c": 3,"d": 4}]'::jsonb);
 test3arr 
----------
        2
(1 row)

-- test jsonb int -> python int
CREATE FUNCTION test1int(val jsonb) RETURNS int
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
assert(val == 1)
return val
$$;
SELECT test1int('1'::jsonb);
 test1int 
----------
        1
(1 row)

-- test jsonb string -> python string
CREATE FUNCTION test1string(val jsonb) RETURNS text
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
assert(val == "a")
return val
$$;
SELECT test1string('"a"'::jsonb);
 test1string 
-------------
 a
(1 row)

-- test jsonb null -> python None
CREATE FUNCTION test1null(val jsonb) RETURNS int
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
assert(val == None)
return 1
$$;
SELECT test1null('null'::jsonb);
 test1null 
-----------
         1
(1 row)

-- test python -> jsonb
CREATE FUNCTION roundtrip(val jsonb) RETURNS jsonb
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
as $$
return val
$$;
SELECT roundtrip('null'::jsonb);
 roundtrip 
-----------
 
(1 row)

SELECT roundtrip('1'::jsonb);
 roundtrip 
-----------
 1
(1 row)

SELECT roundtrip('1234567890.0987654321'::jsonb);
       roundtrip       
-----------------------
 1234567890.0987654321
(1 row)

SELECT roundtrip('-1234567890.0987654321'::jsonb);
       roundtrip        
------------------------
 -1234567890.0987654321
(1 row)

SELECT roundtrip('true'::jsonb);
 roundtrip 
-----------
 true
(1 row)

SELECT roundtrip('"string"'::jsonb);
 roundtrip 
-----------
 "string"
(1 row)

SELECT roundtrip('{"1": null}'::jsonb);
  roundtrip  
-------------
 {"1": null}
(1 row)

SELECT roundtrip('{"1": 1}'::jsonb);
 roundtrip 
-----------
 {"1": 1}
(1 row)

SELECT roundtrip('{"1": true}'::jsonb);
  roundtrip  
-------------
 {"1": true}
(1 row)

SELECT roundtrip('{"1": "string"}'::jsonb);
    roundtrip    
-----------------
 {"1": "string"}
(1 row)

SELECT roundtrip('[null]'::jsonb);
 roundtrip 
-----------
 [null]
(1 row)

SELECT roundtrip('[1]'::jsonb);
 roundtrip 
-----------
 [1]
(1 row)

SELECT roundtrip('[true]'::jsonb);
 roundtrip 
-----------
 [true]
(1 row)

SELECT roundtrip('["string"]'::jsonb);
 roundtrip  
------------
 ["string"]
(1 row)

SELECT roundtrip('[null, 1]'::jsonb);
 roundtrip 
-----------
 [null, 1]
(1 row)

SELECT roundtrip('[1, true]'::jsonb);
 roundtrip 
-----------
 [1, true]
(1 row)

SELECT roundtrip('[true, "string"]'::jsonb);
    roundtrip     
------------------
 [true, "string"]
(1 row)

SELECT roundtrip('["string", "string2"]'::jsonb);
       roundtrip       
-----------------------
 ["string", "string2"]
(1 row)

-- complex numbers -> jsonb
CREATE FUNCTION testComplexNumbers() RETURNS jsonb
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
x = 1 + 2j
return x
$$;
SELECT testComplexNumbers();
ERROR:  could not convert value "(1+2j)" to jsonb
CONTEXT:  while creating return value
PL/Python function "testcomplexnumbers"
-- range -> jsonb
CREATE FUNCTION testRange() RETURNS jsonb
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
x = range(3)
return x
$$;
SELECT testRange();
 testrange 
-----------
 [0, 1, 2]
(1 row)

-- 0xff -> jsonb
CREATE FUNCTION testDecimal() RETURNS jsonb
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
x = 0xff
return x
$$;
SELECT testDecimal();
 testdecimal 
-------------
 255
(1 row)

-- tuple -> jsonb
CREATE FUNCTION testTuple() RETURNS jsonb
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
x = (1, 'String', None)
return x
$$;
SELECT testTuple();
      testtuple      
---------------------
 [1, "String", null]
(1 row)

-- interesting dict -> jsonb
CREATE FUNCTION test_dict1() RETURNS jsonb
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
x = {"a": 1, None: 2, 33: 3}
return x
$$;
SELECT test_dict1();
        test_dict1        
--------------------------
 {"": 2, "a": 1, "33": 3}
(1 row)