1 # Licensed under the Apache License, Version 2.0 (the "License");
2 # you may not use this file except in compliance with the License.
3 # You may obtain a copy of the License at
5 # http://www.apache.org/licenses/LICENSE-2.0
7 # Unless required by applicable law or agreed to in writing, software
8 # distributed under the License is distributed on an "AS IS" BASIS,
9 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10 # See the License for the specific language governing permissions and
11 # limitations under the License.
13 from __future__
import absolute_import
14 from __future__
import division
15 from __future__
import print_function
16 from __future__
import unicode_literals
18 from thrift
.Thrift
import TType
22 SPEC_ARGS_CLASS_REF_IDX
= 0
23 SPEC_ARGS_THRIFT_SPEC_IDX
= 1
26 def fix_spec(all_structs
):
27 """Wire up recursive references for all TStruct definitions inside of each thrift_spec."""
28 for struc
in all_structs
:
29 spec
= struc
.thrift_spec
30 for thrift_spec
in spec
:
31 if thrift_spec
is None:
33 elif thrift_spec
[TYPE_IDX
] == TType
.STRUCT
:
34 other
= thrift_spec
[SPEC_ARGS_IDX
][SPEC_ARGS_CLASS_REF_IDX
].thrift_spec
35 thrift_spec
[SPEC_ARGS_IDX
][SPEC_ARGS_THRIFT_SPEC_IDX
] = other
36 elif thrift_spec
[TYPE_IDX
] in (TType
.LIST
, TType
.SET
):
37 _fix_list_or_set(thrift_spec
[SPEC_ARGS_IDX
])
38 elif thrift_spec
[TYPE_IDX
] == TType
.MAP
:
39 _fix_map(thrift_spec
[SPEC_ARGS_IDX
])
42 def _fix_list_or_set(element_type
):
43 # For a list or set, the thrift_spec entry looks like,
44 # (1, TType.LIST, 'lister', (TType.STRUCT, [RecList, None], False), None, ), # 1
45 # so ``element_type`` will be,
46 # (TType.STRUCT, [RecList, None], False)
47 if element_type
[0] == TType
.STRUCT
:
48 element_type
[1][1] = element_type
[1][0].thrift_spec
49 elif element_type
[0] in (TType
.LIST
, TType
.SET
):
50 _fix_list_or_set(element_type
[1])
51 elif element_type
[0] == TType
.MAP
:
52 _fix_map(element_type
[1])
55 def _fix_map(element_type
):
56 # For a map of key -> value type, ``element_type`` will be,
57 # (TType.I16, None, TType.STRUCT, [RecMapBasic, None], False), None, )
58 # which is just a normal struct definition.
60 # For a map of key -> list / set, ``element_type`` will be,
61 # (TType.I16, None, TType.LIST, (TType.STRUCT, [RecMapList, None], False), False)
62 # and we need to process the 3rd element as a list.
64 # For a map of key -> map, ``element_type`` will be,
65 # (TType.I16, None, TType.MAP, (TType.I16, None, TType.STRUCT,
66 # [RecMapMap, None], False), False)
67 # and need to process 3rd element as a map.
69 # Is the map key a struct?
70 if element_type
[0] == TType
.STRUCT
:
71 element_type
[1][1] = element_type
[1][0].thrift_spec
72 elif element_type
[0] in (TType
.LIST
, TType
.SET
):
73 _fix_list_or_set(element_type
[1])
74 elif element_type
[0] == TType
.MAP
:
75 _fix_map(element_type
[1])
77 # Is the map value a struct?
78 if element_type
[2] == TType
.STRUCT
:
79 element_type
[3][1] = element_type
[3][0].thrift_spec
80 elif element_type
[2] in (TType
.LIST
, TType
.SET
):
81 _fix_list_or_set(element_type
[3])
82 elif element_type
[2] == TType
.MAP
:
83 _fix_map(element_type
[3])