Coverage for intelligence_toolkit/match_entity_records/classes.py: 100%

1# Copyright (c) 2024 Microsoft Corporation. All rights reserved. 

2# Licensed under the MIT license. See LICENSE file in the project. 

3# 

4 

5from typing import TypedDict 

6 

7import polars as pl 

8from pydantic import BaseModel 

9 

10 

11class RecordsModel(BaseModel): 

12 dataframe: pl.DataFrame 

13 name_column: str 

14 columns: list[str] 

15 dataframe_name: str | None = None 

16 id_column: str | None = None 

17 

18 class Config: 

19 arbitrary_types_allowed = True 

20 

21 

22class AttributeToMatch(TypedDict): 

23 label: str | None = None 

24 columns: list[str]