Coverage for tests / test_select_typing.py: 100%

1from sqlmodel import Field, Session, SQLModel, create_engine, select 8 ctx1ijklmnop

2from sqlmodel.pool import StaticPool 8 ctx1ijklmnop

3 

4 

5def test_fields() -> None: 8 ctx1ijklmnop

6 class Hero(SQLModel, table=True): 8 ctx1abcdefgh

7 id: int | None = Field(default=None, primary_key=True) 8 ctx1abcdefgh

8 name: str 8 ctx1abcdefgh

9 secret_name: str 8 ctx1abcdefgh

10 age: int | None = None 8 ctx1abcdefgh

11 food: str | None = None 8 ctx1abcdefgh

12 

13 engine = create_engine( 8 ctx1abcdefgh

14 "sqlite://", connect_args={"check_same_thread": False}, poolclass=StaticPool 

15 ) 

16 

17 SQLModel.metadata.create_all(engine) 8 ctx1abcdefgh

18 

19 with Session(engine) as session: 8 ctx1abcdefgh

20 session.add(Hero(name="Deadpond", secret_name="Dive Wilson")) 8 ctx1abcdefgh

21 session.add( 8 ctx1abcdefgh

22 Hero(name="Spider-Boy", secret_name="Pedro Parqueador", food="pizza") 

23 ) 

24 session.add(Hero(name="Rusty-Man", secret_name="Tommy Sharp", age=48)) 8 ctx1abcdefgh

25 

26 session.commit() 8 ctx1abcdefgh

27 

28 # check typing of select with 3 fields 

29 with Session(engine) as session: 8 ctx1abcdefgh

30 statement_3 = select(Hero.id, Hero.name, Hero.secret_name) 8 ctx1abcdefgh

31 results_3 = session.exec(statement_3) 8 ctx1abcdefgh

32 for hero_3 in results_3: 8 ctx1abcdefgh

33 assert len(hero_3) == 3 8 ctx1abcdefgh

34 name_3: str = hero_3[1] 8 ctx1abcdefgh

35 assert type(name_3) is str 8 ctx1abcdefgh

36 assert type(hero_3[0]) is int 8 ctx1abcdefgh

37 assert type(hero_3[2]) is str 8 ctx1abcdefgh

38 

39 # check typing of select with 4 fields 

40 with Session(engine) as session: 8 ctx1abcdefgh

41 statement_4 = select(Hero.id, Hero.name, Hero.secret_name, Hero.age) 8 ctx1abcdefgh

42 results_4 = session.exec(statement_4) 8 ctx1abcdefgh

43 for hero_4 in results_4: 8 ctx1abcdefgh

44 assert len(hero_4) == 4 8 ctx1abcdefgh

45 name_4: str = hero_4[1] 8 ctx1abcdefgh

46 assert type(name_4) is str 8 ctx1abcdefgh

47 assert type(hero_4[0]) is int 8 ctx1abcdefgh

48 assert type(hero_4[2]) is str 8 ctx1abcdefgh

49 assert type(hero_4[3]) in [int, type(None)] 8 ctx1abcdefgh

50 

51 # check typing of select with 5 fields: currently runs but doesn't pass mypy 

52 # with Session(engine) as session: 

53 # statement_5 = select(Hero.id, Hero.name, Hero.secret_name, Hero.age, Hero.food) 

54 # results_5 = session.exec(statement_5) 

55 # for hero_5 in results_5: 

56 # assert len(hero_5) == 5 

57 # name_5: str = hero_5[1] 

58 # assert type(name_5) is str 

59 # assert type(hero_5[0]) is int 

60 # assert type(hero_5[2]) is str 

61 # assert type(hero_5[3]) in [int, type(None)] 

62 # assert type(hero_5[4]) in [str, type(None)]