def test_valid_map():
tmap = TMap(key_type=TI32(), value_type=TString())
assert tmap.validate_arg({}) is None
assert tmap.validate_arg({3: "dog", 4: "cat"}) is None
nested_tmap = TMap(key_type=TI32(),
value_type=TMap(key_type=TString(), value_type=TI32()))
assert nested_tmap.validate_arg({3: {"test": 4}, 1: {"map": 3231}}) is None
def _parse_type(type_info):
try:
ttype_code, nested_type_info = type_info
except TypeError:
ttype_code = type_info
nested_type_info = None
ttype = TType._VALUES_TO_NAMES[ttype_code].lower()
if nested_type_info is None:
return _BASIC_TYPE_MAP[ttype]()
elif ttype == "list":
return TList(value_type=_parse_type(nested_type_info))
elif ttype == "set":
return TSet(value_type=_parse_type(nested_type_info))
elif ttype == "map":
key, value = nested_type_info
return TMap(key_type=_parse_type(key), value_type=_parse_type(value))
elif ttype == "struct":
return TStruct(
name=nested_type_info.__name__,
fields=[
_parse_arg(result)
for result in nested_type_info.thrift_spec.values()
],
)
# Its a basic type but has defined nested type info. its probably an enum
return TEnum(
name=nested_type_info.__name__,
names_to_values=nested_type_info._NAMES_TO_VALUES,
values_to_names=nested_type_info._VALUES_TO_NAMES,
)
def test_invalid_map():
tmap = TMap(key_type=TI32(), value_type=TString())
assert tmap.validate_arg(4) == ["Expected dict but got int"]
assert tmap.validate_arg({
3: 4,
5: 2
}) == [
"Value for key '3' in map invalid: 'Expected str but got int'",
"Value for key '5' in map invalid: 'Expected str but got int'",
]
assert tmap.validate_arg({
"3": 4,
5: "2"
}) == [
"Key '3' in map invalid: 'Expected int but got str'",
"Value for key '3' in map invalid: 'Expected str but got int'",
]
def test_super_nesting():
turducken_thrift = load_thrift_from_testdir("turducken.thrift")
dog_names_to_values = {"GOLDEN": 0, "CORGI": 1, "BASSET": 2}
dog_enum = TEnum(
name="DOG",
names_to_values=dog_names_to_values,
values_to_names={
value: key
for key, value in dog_names_to_values.items()
},
)
outer_value = TI64()
inner_value = TList(value_type=TString())
inner_key = TSet(value_type=TList(value_type=dog_enum))
outer_key = TMap(key_type=inner_key, value_type=inner_value)
map_field = TMap(key_type=outer_key, value_type=outer_value)
expected = ServiceEndpoint(
name="getTheStruct",
args=[],
results=[
ThriftSpec(
field_id=0,
name="success",
type_info=TStruct(
name="TheStruct",
fields=[
ThriftSpec(
field_id=1,
name="myInsaneStruct",
type_info=map_field,
required=True,
)
],
),
required=False,
)
],
)
assert expected == thrift_parser._parse_thrift_endpoint(
turducken_thrift.__thrift_meta__["services"][0], "getTheStruct")
def test_map_type():
collections_thrift = load_thrift_from_testdir("collections.thrift")
expected = ServiceEndpoint(
name="maps",
args=[
ThriftSpec(
name="mapofI16toI64",
type_info=TMap(key_type=TI16(), value_type=TI64()),
required=False,
)
],
results=[
ThriftSpec(
name="success",
type_info=TMap(key_type=TBool(), value_type=TByte()),
required=False,
)
],
)
result = thrift_parser._parse_thrift_endpoint(
collections_thrift.__thrift_meta__["services"][0], "maps")
assert expected == result