def _test_schema(pb_msg_type, expected_schema, custom_fields=None):
field_map = {field.name: field for field in expected_schema.fields}
fields = pb_msg_type.DESCRIPTOR.fields_by_name
for field_name, field_desc in fields.iteritems():
field, _ = pb2df.convert_field(field_desc, custom_fields)
assert field == field_map[field_name]
schema, _ = pb2df.convert_schema(pb_msg_type.DESCRIPTOR, custom_fields)
assert schema == expected_schema
def test_custom_field_getter(custom_field_msg):
convert = datetime.utcfromtimestamp
custom_fields = {
'timestamp_field': pb2df.CustomField(
types.TimestampType(), convert),
}
pb_msg_type = custom_field_msg.__class__
field_desc = pb_msg_type.DESCRIPTOR.fields_by_name['timestamp_field']
_, getter = pb2df.convert_field(field_desc, custom_fields)
assert getter(custom_field_msg) == convert(custom_field_msg.timestamp_field)
def test_repeated_nested_field_getter(nested_msg):
pb_msg_type = nested_msg.__class__
field_desc = pb_msg_type.DESCRIPTOR.fields_by_name['repeated_nested_field']
_, getter = pb2df.convert_field(field_desc)
assert getter(nested_msg) == [(1,), (2,)]
def test_nested_field_getter(nested_msg):
pb_msg_type = nested_msg.__class__
field_desc = pb_msg_type.DESCRIPTOR.fields_by_name['required_nested_field']
_, getter = pb2df.convert_field(field_desc)
assert getter(nested_msg) == (999,)
def test_empty_field_getter():
pb_msg_type = example_pb2.SimpleMessage
field_desc = pb_msg_type.DESCRIPTOR.fields_by_name['field']
_, getter = pb2df.convert_field(field_desc)
assert getter(pb_msg_type()) is None
def test_default_field_getter():
pb_msg_type = example_pb2.LabeledMessage
field_desc = pb_msg_type.DESCRIPTOR.fields_by_name['default_field']
_, getter = pb2df.convert_field(field_desc)
assert getter(pb_msg_type()) == field_desc.default_value
def test_repeated_field_getter(labeled_msg):
pb_msg_type = labeled_msg.__class__
field_desc = pb_msg_type.DESCRIPTOR.fields_by_name['repeated_field']
_, getter = pb2df.convert_field(field_desc)
assert getter(labeled_msg) == labeled_msg.repeated_field
def test_basic_field_getter(basic_msg):
pb_msg_type = basic_msg.__class__
field_desc = pb_msg_type.DESCRIPTOR.fields_by_name['double_field']
_, getter = pb2df.convert_field(field_desc)
assert getter(basic_msg) == basic_msg.double_field