Coverage for tests / test_select_typing.py: 100%

1from sqlalchemy import Engine 9 ctx1BCDEFGHIJ

2from sqlalchemy.engine import ScalarResult, TupleResult 9 ctx1BCDEFGHIJ

3from sqlalchemy.sql.elements import UnaryExpression 9 ctx1BCDEFGHIJ

4from sqlmodel import ( 9 ctx1BCDEFGHIJ

5 Field, 

6 Session, 

7 SQLModel, 

8 asc, 

9 col, 

10 desc, 

11 nulls_first, 

12 nulls_last, 

13 select, 

14) 

15from typing_extensions import assert_type 9 ctx1BCDEFGHIJ

16 

17 

18def test_fields(database_engine: Engine) -> None: 9 ctx1BCDEFGHIJ

19 class Hero(SQLModel, table=True): 9 ctx1abcdefghi

20 id: int | None = Field(default=None, primary_key=True) 9 ctx1abcdefghi

21 name: str 9 ctx1abcdefghi

22 secret_name: str 9 ctx1abcdefghi

23 age: int | None = None 9 ctx1abcdefghi

24 food: str | None = None 9 ctx1abcdefghi

25 

26 engine = database_engine 9 ctx1abcdefghi

27 

28 SQLModel.metadata.create_all(engine) 9 ctx1abcdefghi

29 

30 with Session(engine) as session: 9 ctx1abcdefghi

31 session.add(Hero(name="Deadpond", secret_name="Dive Wilson")) 9 ctx1abcdefghi

32 session.add( 9 ctx1abcdefghi

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

34 ) 

35 session.add(Hero(name="Rusty-Man", secret_name="Tommy Sharp", age=48)) 9 ctx1abcdefghi

36 

37 session.commit() 9 ctx1abcdefghi

38 

39 # check typing of select with 3 fields 

40 with Session(engine) as session: 9 ctx1abcdefghi

41 statement_3 = select(col(Hero.id), Hero.name, Hero.secret_name) 9 ctx1abcdefghi

42 results_3 = session.exec(statement_3) 9 ctx1abcdefghi

43 assert_type(results_3, TupleResult[tuple[int | None, str, str]]) 9 ctx1abcdefghi

44 for hero_3 in results_3: 9 ctx1abcdefghi

45 assert len(hero_3) == 3 9 ctx1abcdefghi

46 name_3: str = hero_3[1] 9 ctx1abcdefghi

47 assert type(name_3) is str 9 ctx1abcdefghi

48 assert type(hero_3[0]) is int 9 ctx1abcdefghi

49 assert type(hero_3[2]) is str 9 ctx1abcdefghi

50 

51 # check typing of select with 4 fields 

52 with Session(engine) as session: 9 ctx1abcdefghi

53 statement_4 = select(col(Hero.id), Hero.name, Hero.secret_name, col(Hero.age)) 9 ctx1abcdefghi

54 results_4 = session.exec(statement_4) 9 ctx1abcdefghi

55 assert_type(results_4, TupleResult[tuple[int | None, str, str, int | None]]) 9 ctx1abcdefghi

56 for hero_4 in results_4: 9 ctx1abcdefghi

57 assert len(hero_4) == 4 9 ctx1abcdefghi

58 name_4: str = hero_4[1] 9 ctx1abcdefghi

59 assert type(name_4) is str 9 ctx1abcdefghi

60 assert type(hero_4[0]) is int 9 ctx1abcdefghi

61 assert type(hero_4[2]) is str 9 ctx1abcdefghi

62 assert type(hero_4[3]) in [int, type(None)] 9 ctx1abcdefghi

63 

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

65 # with Session(engine) as session: 

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

67 # results_5 = session.exec(statement_5) 

68 # for hero_5 in results_5: 

69 # assert len(hero_5) == 5 

70 # name_5: str = hero_5[1] 

71 # assert type(name_5) is str 

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

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

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

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

76 

77 

78def test_scalar_results(database_engine: Engine) -> None: 9 ctx1BCDEFGHIJ

79 class Hero(SQLModel, table=True): 9 ctx1jklmnopqr

80 id: int | None = Field(default=None, primary_key=True) 9 ctx1jklmnopqr

81 name: str 9 ctx1jklmnopqr

82 

83 SQLModel.metadata.create_all(database_engine) 9 ctx1jklmnopqr

84 with Session(database_engine) as session: 9 ctx1jklmnopqr

85 session.add(Hero(name="Deadpond")) 9 ctx1jklmnopqr

86 session.commit() 9 ctx1jklmnopqr

87 

88 statement = select(Hero) 9 ctx1jklmnopqr

89 result = session.exec(statement) 9 ctx1jklmnopqr

90 assert_type(result, ScalarResult[Hero]) 9 ctx1jklmnopqr

91 assert result.one().name == "Deadpond" 9 ctx1jklmnopqr

92 

93 scalars = session.scalars(statement) 9 ctx1jklmnopqr

94 assert_type(scalars, ScalarResult[Hero]) 9 ctx1jklmnopqr

95 assert scalars.one().name == "Deadpond" 9 ctx1jklmnopqr

96 

97 scalar_subquery = ( 9 ctx1jklmnopqr

98 select(Hero.name) 

99 .where(Hero.name == "Deadpond") 

100 .group_by(Hero.name) 

101 .having(Hero.name == "Deadpond") 

102 .scalar_subquery() 

103 ) 

104 name = session.exec(select(scalar_subquery)).one() 9 ctx1jklmnopqr

105 assert_type(name, str) 9 ctx1jklmnopqr

106 assert name == "Deadpond" 9 ctx1jklmnopqr

107 

108 

109def test_ordering_types() -> None: 9 ctx1BCDEFGHIJ

110 class Hero(SQLModel, table=True): 9 ctx1stuvwxyzA

111 id: int = Field(primary_key=True) 9 ctx1stuvwxyzA

112 name: str 9 ctx1stuvwxyzA

113 

114 ascending = asc(Hero.id) 9 ctx1stuvwxyzA

115 assert_type(ascending, UnaryExpression[int]) 9 ctx1stuvwxyzA

116 assert str(ascending) == "hero.id ASC" 9 ctx1stuvwxyzA

117 

118 descending = desc(Hero.id) 9 ctx1stuvwxyzA

119 assert_type(descending, UnaryExpression[int]) 9 ctx1stuvwxyzA

120 assert str(descending) == "hero.id DESC" 9 ctx1stuvwxyzA

121 

122 first = nulls_first(Hero.name) 9 ctx1stuvwxyzA

123 assert_type(first, UnaryExpression[str]) 9 ctx1stuvwxyzA

124 assert str(first) == "hero.name NULLS FIRST" 9 ctx1stuvwxyzA

125 

126 last = nulls_last(Hero.name) 9 ctx1stuvwxyzA

127 assert_type(last, UnaryExpression[str]) 9 ctx1stuvwxyzA

128 assert str(last) == "hero.name NULLS LAST" 9 ctx1stuvwxyzA