SNOW-1794359:[API Coverage] MapType (#2636)

snowflakedb · Nov 18, 2024 · d4b03af · d4b03af
1 parent 8976902
commit d4b03af
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@
   - `simple_string`
   - `json_value`
   - `json`
+- Added support for variables `keyType` and `valueType` in class `MapType`
 
 #### Improvements
 

diff --git a/src/snowflake/snowpark/types.py b/src/snowflake/snowpark/types.py
@@ -328,6 +328,14 @@ def json_value(self) -> Dict[str, Any]:
             "value_type": self.value_type.json_value(),
         }
 
+    @property
+    def keyType(self):
+        return self.key_type
+
+    @property
+    def valueType(self):
+        return self.value_type
+
     simpleString = simple_string
     jsonValue = json_value
 

diff --git a/tests/unit/test_types.py b/tests/unit/test_types.py
@@ -1234,3 +1234,14 @@ def test_datatype(tpe, simple_string, json, type_name, json_value):
             tpe.typeName()
     else:
         assert tpe.typeName() == type_name
+
+
+def test_maptype_alias():
+    expected_key = StringType()
+    expected_value = IntegerType()
+    tpe = MapType(expected_key, expected_value)
+    assert tpe.valueType == expected_value
+    assert tpe.keyType == expected_key
+
+    assert tpe.valueType == tpe.value_type
+    assert tpe.keyType == tpe.key_type