| 160 | |
| 161 | |
| 162 | class ForyTypeVisitor(TypeVisitor): |
| 163 | def visit_array(self, field_name, elem_type, carrier, types_path=None): |
| 164 | raise TypeError(f"Row format does not support pyfory.{carrier} array annotations") |
| 165 | |
| 166 | def visit_list(self, field_name, elem_type, types_path=None): |
| 167 | # Infer type recursively for type such as List[Dict[str, str]] |
| 168 | elem_field = infer_field("item", elem_type, self, types_path=types_path) |
| 169 | return field(field_name, list_(elem_field.type)) |
| 170 | |
| 171 | def visit_set(self, field_name, elem_type, types_path=None): |
| 172 | # Infer type recursively for type such as Set[Dict[str, str]] |
| 173 | elem_field = infer_field("item", elem_type, self, types_path=types_path) |
| 174 | return field(field_name, list_(elem_field.type)) |
| 175 | |
| 176 | def visit_tuple(self, field_name, elem_types, types_path=None): |
| 177 | elem_type = get_homogeneous_tuple_elem_type(elem_types) |
| 178 | if elem_type is None: |
| 179 | raise TypeError(f"Row format supports only homogeneous tuple annotations, got {elem_types}") |
| 180 | elem_field = infer_field("item", elem_type, self, types_path=types_path) |
| 181 | return field(field_name, list_(elem_field.type)) |
| 182 | |
| 183 | def visit_dict(self, field_name, key_type, value_type, types_path=None): |
| 184 | # Infer type recursively for type such as Dict[str, Dict[str, str]] |
| 185 | key_field = infer_field("key", key_type, self, types_path=types_path) |
| 186 | value_field = infer_field("value", value_type, self, types_path=types_path) |
| 187 | return field(field_name, map_(key_field.type, value_field.type)) |
| 188 | |
| 189 | def visit_customized(self, field_name, type_, types_path=None): |
| 190 | # type_ is a pojo |
| 191 | pojo_schema = infer_schema(type_) |
| 192 | fields = [pojo_schema.field(i) for i in range(pojo_schema.num_fields)] |
| 193 | # TODO: Add metadata support |
| 194 | return field(field_name, struct(fields)) |
| 195 | |
| 196 | def visit_other(self, field_name, type_, types_path=None): |
| 197 | if type_ in _supported_types_mapping: |
| 198 | fory_type_func = _supported_types_mapping.get(type_) |
| 199 | return field(field_name, fory_type_func()) |
| 200 | if isinstance(type_, type) and type_.__module__ != "builtins": |
| 201 | return self.visit_customized(field_name, type_, types_path=types_path) |
| 202 | raise TypeError( |
| 203 | f"Type {type_} not supported, currently only compositions of {_supported_types_str} are supported. types_path is {types_path}" |
| 204 | ) |
| 205 | |
| 206 | |
| 207 | def infer_data_type(clz) -> Optional[DataType]: |
no outgoing calls