File: pg_stuff.py

package info (click to toggle)
sqlalchemy 2.0.40%2Bds1-1
  • links: PTS, VCS
  • area: main
  • in suites: trixie
  • size: 26,404 kB
  • sloc: python: 410,002; makefile: 230; sh: 7
file content (133 lines) | stat: -rw-r--r-- 3,505 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
from typing import Any
from typing import Dict
from uuid import UUID as _py_uuid

from sqlalchemy import cast
from sqlalchemy import Column
from sqlalchemy import func
from sqlalchemy import Integer
from sqlalchemy import or_
from sqlalchemy import select
from sqlalchemy import Text
from sqlalchemy import UniqueConstraint
from sqlalchemy.dialects.postgresql import ARRAY
from sqlalchemy.dialects.postgresql import array
from sqlalchemy.dialects.postgresql import DATERANGE
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.dialects.postgresql import INT4RANGE
from sqlalchemy.dialects.postgresql import INT8MULTIRANGE
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.dialects.postgresql import TSTZMULTIRANGE
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column

# test #6402

c1 = Column(UUID())

# EXPECTED_TYPE: Column[UUID]
reveal_type(c1)

c2 = Column(UUID(as_uuid=False))

# EXPECTED_TYPE: Column[str]
reveal_type(c2)


class Base(DeclarativeBase):
    pass


class Test(Base):
    __tablename__ = "test_table_json"

    id = mapped_column(Integer, primary_key=True)
    data: Mapped[Dict[str, Any]] = mapped_column(JSONB)

    ident: Mapped[_py_uuid] = mapped_column(UUID())

    ident_str: Mapped[str] = mapped_column(UUID(as_uuid=False))


elem = func.jsonb_array_elements(Test.data, type_=JSONB).column_valued("elem")

stmt = select(Test).where(
    or_(
        cast("example code", ARRAY(Text)).contained_by(
            array([select(elem["code"].astext).scalar_subquery()])
        ),
        cast("stefan", ARRAY(Text)).contained_by(
            array([select(elem["code"]["new_value"].astext).scalar_subquery()])
        ),
    )
)
print(stmt)


t1 = Test()

# EXPECTED_RE_TYPE: .*[dD]ict\[.*str, Any\]
reveal_type(t1.data)

# EXPECTED_TYPE: UUID
reveal_type(t1.ident)

unique = UniqueConstraint(name="my_constraint")
insert(Test).on_conflict_do_nothing(
    "foo", [Test.id], Test.id > 0
).on_conflict_do_update(
    unique, ["foo"], Test.id > 0, {"id": 42, Test.ident: 99}, Test.id == 22
).excluded.foo.desc()

s1 = insert(Test)
s1.on_conflict_do_update(set_=s1.excluded)


# EXPECTED_TYPE: Column[Range[int]]
reveal_type(Column(INT4RANGE()))
# EXPECTED_TYPE: Column[Range[datetime.date]]
reveal_type(Column("foo", DATERANGE()))
# EXPECTED_TYPE: Column[Sequence[Range[int]]]
reveal_type(Column(INT8MULTIRANGE()))
# EXPECTED_TYPE: Column[Sequence[Range[datetime.datetime]]]
reveal_type(Column("foo", TSTZMULTIRANGE()))


range_col_stmt = select(Column(INT4RANGE()), Column(INT8MULTIRANGE()))

# EXPECTED_TYPE: Select[Tuple[Range[int], Sequence[Range[int]]]]
reveal_type(range_col_stmt)

array_from_ints = array(range(2))

# EXPECTED_TYPE: array[int]
reveal_type(array_from_ints)

array_of_strings = array([], type_=Text)

# EXPECTED_TYPE: array[str]
reveal_type(array_of_strings)

array_of_ints = array([0], type_=Integer)

# EXPECTED_TYPE: array[int]
reveal_type(array_of_ints)

# EXPECTED_MYPY: Cannot infer type argument 1 of "array"
array([0], type_=Text)

# EXPECTED_TYPE: ARRAY[str]
reveal_type(ARRAY(Text))

# EXPECTED_TYPE: Column[Sequence[int]]
reveal_type(Column(type_=ARRAY(Integer)))

stmt_array_agg = select(func.array_agg(Column("num", type_=Integer)))

# EXPECTED_TYPE: Select[Tuple[Sequence[int]]]
reveal_type(stmt_array_agg)

# EXPECTED_TYPE: Select[Tuple[Sequence[str]]]
reveal_type(select(func.array_agg(Test.ident_str)))