def test_get_record() -> None:
data_table = DataTable({"a": [1, 2], "b": [3, 4]})
assert data_table.get_record(0) == {"a": 1, "b": 3}
assert data_table.get_record(1) == {"a": 2, "b": 4}
with pytest.raises(DataTableError):
data_table.get_record(2)
with pytest.raises(DataTableError):
DataTable({"a": [1, 2], "b": [3]}).get_record(2)
def test_drop_duplicates(self):
class MyRecord(DictClass):
first_name: str
last_name: str
sex: str
data_table = DataTable(
record_class=MyRecord,
base_dict={
"first_name": ["ABC", "ABC", "DEF", "DEF", "DEF"],
"last_name": ["XYZ", "XYZ", "MNO", "MNO", "MNO"],
"sex": ["M", "M", "F", "M", "F"],
},
)
assert isinstance(data_table.get_record(0), MyRecord)
deduplicated_data_table = data_table.drop_duplicates()
assert isinstance(deduplicated_data_table.get_record(0), MyRecord)
assert {
"first_name": ["ABC", "DEF", "DEF"],
"last_name": ["XYZ", "MNO", "MNO"],
"sex": ["M", "F", "M"],
} == deduplicated_data_table
deduplicated_data_table = data_table.drop_duplicates(
subset=("first_name", ))
assert isinstance(deduplicated_data_table.get_record(0), MyRecord)
assert {
"first_name": ["ABC", "DEF"],
"last_name": ["XYZ", "MNO"],
"sex": ["M", "F"],
} == deduplicated_data_table
deduplicated_data_table = data_table.drop_duplicates(
subset=("last_name", "sex"))
assert isinstance(deduplicated_data_table.get_record(0), MyRecord)
assert {
"first_name": ["ABC", "DEF", "DEF"],
"last_name": ["XYZ", "MNO", "MNO"],
"sex": ["M", "F", "M"],
} == deduplicated_data_table
# it should throw an error as the column name provied is invalid
with pytest.raises(DataTableError):
_ = data_table.drop_duplicates(subset=("invalid_column", ))
# not normalized table
data_table = DataTable(
record_class=MyRecord,
base_dict={
"first_name": ["ABC", "ABC", "DEF", "DEF", "DEF"],
"last_name": ["XYZ", "XYZ", "MNO", "MNO"],
"sex": ["M", "M", "F", "M", "F"],
},
)
# It should throw an error as table is not normalized
with pytest.raises(DataTableError):
_ = data_table.drop_duplicates()
def test_normalize() -> None:
data_table = DataTable({"a": [1, 2, 3], "b": [3, 4], "c": []})
data_table.normalize()
assert data_table == {
"a": [1, 2, 3],
"b": [3, 4, data_table.NOT_SET],
"c": [data_table.NOT_SET, data_table.NOT_SET, data_table.NOT_SET],
}
assert data_table.get_record(0) == {"a": 1, "b": 3, "c": None}
data_table = DataTable({"a": [1, 2], "b": [3, 4]})
data_table.normalize()
assert data_table == {"a": [1, 2], "b": [3, 4]}
def test_custom_record(self):
data_table = DataTable(record_class=UserRecord)
data_table.add_record({"name": "Jon"})
data_table.add_record(UserRecord(name="test", age=12))
assert isinstance(data_table.get_record(0), UserRecord)
assert list(data_table.get_records()) == [
UserRecord(name="Jon"),
UserRecord(name="test", age=12),
]
with pytest.raises(ValueError):
data_table.add_record({"unknown": "Jon"})
with pytest.raises(ValueError):
data_table.add_record({})
def main() -> None:
users_table = DataTable(record_class=UserRecord)
users_table.add_record(
{
"email": "*****@*****.**",
"name": "John",
"age": 34
},
{
"email": "*****@*****.**",
"company": "CiscoSystems",
"name": "Mary",
"age": 34
},
)
print("Get John's record:")
print(users_table.get_record(0))
print("All records as a list:")
print(list(users_table.get_records()))
print("Find record with name=Mary:")
print(users_table.filter_records({"name": "Mary"}).get_record(0))