11from __future__ import annotations
22
3- from typing import TYPE_CHECKING , Literal , NoReturn , Protocol , TypedDict , runtime_checkable
3+ from typing import TYPE_CHECKING , Literal , Protocol , TypedDict , TypeVar , runtime_checkable
44
55if TYPE_CHECKING :
66 from collections .abc import Iterable , Iterator , Sequence
77
8- from typing_extensions import NotRequired , TypeIs
8+ from typing_extensions import TypeIs
99
1010__all__ = [
1111 "CompiledSql" ,
1212 "IntoInterpolation" ,
13- "IntoParam" ,
1413 "IntoTemplate" ,
1514 "Param" ,
1615 "SupportsDuckdbTemplate" ,
@@ -33,32 +32,36 @@ class Param(TypedDict):
3332 name : str
3433
3534
36- class IntoParam (TypedDict ):
37- """A Param with a name that is None, which can be used as input to the template engine, which will assign it a name based on its position and optionally an expression."""
35+ # class IntoParam(TypedDict):
36+ # """A Param with a name that is None, which can be used as input to the template engine, which will assign it a name based on its position and optionally an expression."""
3837
39- value : object
40- name : NotRequired [str | None ]
38+ # value: object
39+ # name: NotRequired[str | None]
40+
41+
42+ # def is_into_param(thing: object) -> TypeIs[IntoParam]:
43+ # try:
44+ # value = thing["value"] # ty:ignore[not-subscriptable]
45+ # except (TypeError, KeyError):
46+ # return False
47+ # try:
48+ # name = thing["name"] # ty:ignore[not-subscriptable]
49+ # except KeyError:
50+ # name = None
4151
52+ # return isinstance(value, object) and isinstance(name, (str, type(None)))
4253
43- def is_into_param (thing : object ) -> TypeIs [IntoParam ]:
44- try :
45- value = thing ["value" ] # ty:ignore[not-subscriptable]
46- except (TypeError , KeyError ):
47- return False
48- try :
49- name = thing ["name" ] # ty:ignore[not-subscriptable]
50- except KeyError :
51- name = None
5254
53- return isinstance (value , object ) and isinstance (name , (str , type (None )))
55+ def is_into_interpolation (thing : object ) -> TypeIs [IntoInterpolation ]:
56+ return isinstance (thing , IntoInterpolation )
5457
5558
5659@runtime_checkable
5760class SupportsDuckdbTemplate (Protocol ):
58- """Something that can be converted into a SqlTemplate by implementing the __duckdb_template__ method."""
61+ """Something that can be converted into a Template by implementing the __duckdb_template__ method."""
5962
60- def __duckdb_template__ (self , / , ** future_kwargs ) -> str | IntoParam | Iterable [str | IntoParam ]:
61- """Convert self into a SqlTemplate , by returning either a string, an IntoParam , or an iterable of these.
63+ def __duckdb_template__ (self , / , ** future_kwargs ) -> str | IntoInterpolation | Iterable [str | IntoInterpolation ]:
64+ """Convert self into an IntoTemplate , by returning either a string, an IntoInterpolation , or an iterable of these.
6265
6366 The future_kwargs are for future extensibility, in case duckdb wants
6467 to pass additional information in the future.
@@ -83,10 +86,8 @@ def __init__(self, user_id: int):
8386 self.user_id = user_id
8487
8588 def __duckdb_template__(self, **kwargs):
86- return [
87- "SELECT * FROM users WHERE id = ",
88- {"value": self.user_id, "name": "user_id"},
89- ]
89+ user_id = self.user_id
90+ return t"SELECT * FROM users WHERE id = {user_id}"
9091 ```
9192
9293 This will resolve to the final SQL and params:
@@ -99,15 +100,15 @@ def __duckdb_template__(self, **kwargs):
99100 """
100101
101102
102- def param (value : object , name : str | None = None ) -> IntoParam :
103+ def param (value : object , name : str | None = None ) -> ParamInterpolation :
103104 """Helper function to create an IntoParam with an optional name."""
104105 if name is not None :
105106 assert_param_name_legal (name )
106- return IntoParam (value = value , name = name )
107+ return ParamInterpolation (value = value , expression = name , conversion = None )
107108
108109
109- def template (thing : object , / , ** ignored_kwargs ) -> SqlTemplate :
110- """Convert something to a SqlTemplate .
110+ def template (thing : object , / , ** ignored_kwargs ) -> OurTemplate :
111+ """Convert something to a Template-ish .
111112
112113 The rules are:
113114 - If the thing has a __duckdb_template__ method, call it and convert the
@@ -121,35 +122,35 @@ def template(thing: object, /, **ignored_kwargs) -> SqlTemplate:
121122 """
122123 if isinstance (thing , SupportsDuckdbTemplate ):
123124 raw = thing .__duckdb_template__ (** ignored_kwargs )
124- return SqlTemplate (raw )
125+ parts = [raw ] if isinstance (raw , str ) or is_into_interpolation (raw ) else list (raw )
126+ return OurTemplate (* parts )
125127 if isinstance (thing , str ):
126- return SqlTemplate (thing )
128+ return OurTemplate (thing )
127129 if isinstance (thing , IntoTemplate ):
128- return resolve_into_template ( thing )
130+ return OurTemplate ( * thing )
129131 if isinstance (thing , IntoInterpolation ):
130- return resolve_interpolation (thing )
131- return SqlTemplate (param (value = thing ))
132+ return OurTemplate (thing )
133+ return OurTemplate (param (value = thing ))
132134
133135
134136def compile (thing : object ) -> CompiledSql :
135- resolved_template = template (thing )
136- return compile_sql_template (resolved_template )
137+ t = template (thing )
138+ resolved = resolve (t )
139+ return compile_parts (resolved )
137140
138141
139- def resolve_into_template ( template : IntoTemplate ) -> SqlTemplate :
140- """Resolve a Template, recursively resolving interpolations and flattening nested templates ."""
141- parts : list [str | IntoParam ] = []
142- for part in template :
142+ def resolve ( parts : Iterable [ str | IntoInterpolation ] ) -> OurTemplate [ str | IntoInterpolation ] :
143+ """Resolve an OurTemplate by recursively resolving any inner templates and interpolations ."""
144+ resolved : list [str | IntoInterpolation ] = []
145+ for part in parts :
143146 if isinstance (part , str ):
144- parts .append (part )
147+ resolved .append (part )
145148 else :
146- inner_parts = resolve_interpolation (part )
147- parts .extend (inner_parts )
148- return SqlTemplate (* parts )
149+ resolved .extend (resolve_interpolation (part ))
150+ return OurTemplate (* resolved )
149151
150152
151- def resolve_interpolation (interp : IntoInterpolation ) -> SqlTemplate :
152- """Resolve something that can be converted into an Interpolation, recursively resolving any inner templates."""
153+ def resolve_interpolation (interp : IntoInterpolation ) -> Iterable [str | IntoInterpolation ]:
153154 value = interp .value
154155 # if conversion specified (!s, !r, !a), treat as raw sql, eg
155156 # name = "Alice"
@@ -163,16 +164,17 @@ def resolve_interpolation(interp: IntoInterpolation) -> SqlTemplate:
163164 # Note that this is potentially unsafe if the value comes from an untrusted source,
164165 # since it could lead to SQL injection vulnerabilities, so it should be used with caution.
165166 if interp .conversion == "s" :
166- return SqlTemplate (str (value ))
167+ return OurTemplate (str (value ))
167168 elif interp .conversion == "r" :
168- return SqlTemplate (repr (value ))
169+ return OurTemplate (repr (value ))
169170 elif interp .conversion == "a" :
170- return SqlTemplate (ascii (value ))
171+ return OurTemplate (ascii (value ))
171172
172173 if isinstance (value , str ):
173174 # do NOT pass to template, since that would treat it as a raw SQL.
174- return SqlTemplate (param (value , name = interp .expression ))
175+ return OurTemplate (param (value , name = interp .expression ))
175176 templ = template (value )
177+ # resolved = resolve(templ)
176178 # If the resolved inner is just a single Interpolation, then just return
177179 # the original value so that we preserve the expression name.
178180 # For example, if we have
@@ -185,8 +187,8 @@ def resolve_interpolation(interp: IntoInterpolation) -> SqlTemplate:
185187 # "SELECT * FROM people WHERE age = $age", with a param $age=37,
186188 # eg with a friendly param name, rather than
187189 # "SELECT * FROM people WHERE age = $p0", with a param $p0=37
188- if len (templ .strings ) == 2 and templ .strings [0 ] == "" and templ .strings [1 ] == "" and len (templ .params ) == 1 :
189- return SqlTemplate (param (value = value , name = interp .expression ))
190+ if len (templ .strings ) == 2 and templ .strings [0 ] == "" and templ .strings [1 ] == "" and len (templ .interpolations ) == 1 :
191+ return OurTemplate (param (value = value , name = interp .expression ))
190192 else :
191193 # We got something nested, eg
192194 # age = 37
@@ -197,65 +199,33 @@ def resolve_interpolation(interp: IntoInterpolation) -> SqlTemplate:
197199 return templ
198200
199201
200- # class DuckdDbPyRelation:
201- # def __duckdb_template__(self, /, **future_kwargs) -> str:
202- # # this would just return the existing SQL.
203- # return "SELECT * FROM some_table"
204-
205-
206- class SqlTemplate :
207- """Very similar to string.templatelib.Template, but instead of Interpolations, we use IntoParams."""
208-
209- def __init__ (self , thing : str | IntoParam | Iterable [str | IntoParam ]) -> None :
210- parts = [thing ] if isinstance (thing , str ) or is_into_param (thing ) else list (thing )
211- strings , params = parse_strings_and_params (parts )
212- for param in params :
213- if name := param .get ("name" ):
214- assert_param_name_legal (name )
215- self .strings = tuple (strings )
216- self .params = tuple (params )
217-
218- def __iter__ (self ) -> Iterator [str | IntoParam ]:
219- for s , i in zip (self .strings , self .params , strict = False ):
220- yield s
221- yield i
222- yield self .strings [- 1 ]
223-
224- def __str__ (self ) -> NoReturn :
225- msg = f"{ self .__class__ .__name__ } cannot be directly converted to string. It needs to be processed by the SQL engine to produce the final SQL string." # noqa: E501
226- raise NotImplementedError (msg )
227-
228- def compile (self ) -> CompiledSql :
229- return compile_sql_template (self )
230-
231-
232- def compile_sql_template (template : SqlTemplate ) -> CompiledSql :
202+ def compile_parts (parts : Iterable [str | IntoInterpolation ], / ) -> CompiledSql :
233203 """Compile a resolved SqlTemplate into a final SQL string with named parameter placeholders, and a list of Params."""
234204 sql_parts : list [str ] = []
235205 params : list [Param ] = []
236- for part in template :
206+ for part in parts :
237207 if isinstance (part , str ):
238208 sql_parts .append (part )
239209 else :
240210 param_name = f"p{ len (params )} "
241- if passed_name := part .get ( "name" ) :
211+ if passed_name := part .expression :
242212 param_name += f"_{ passed_name } "
243213 sql_parts .append (f"${ param_name } " )
244- params .append ({"name" : param_name , "value" : part [ " value" ] })
214+ params .append ({"name" : param_name , "value" : part . value })
245215 return {
246216 "sql" : "" .join (sql_parts ),
247217 "params" : params ,
248218 }
249219
250220
251- # from string.templatelib import Interpolation, Template
252221@runtime_checkable
253222class IntoInterpolation (Protocol ):
254223 """Something that can be converted into a string.templatelib.Interpolation."""
255224
256225 value : object
257226 conversion : Literal ["s" , "r" , "a" ] | None
258227 expression : str | None
228+ format_spec : str
259229
260230
261231@runtime_checkable
@@ -276,9 +246,50 @@ def assert_param_name_legal(name: str) -> None:
276246 # or perhaps we shouldn't even check, just pass it to duckdb and let it error if it's illegal
277247
278248
249+ class ParamInterpolation :
250+ """A simple implementation of IntoInterpolation, for testing purposes."""
251+
252+ def __init__ (
253+ self , value : object , conversion : Literal ["s" , "r" , "a" ] | None = None , expression : str | None = None
254+ ) -> None :
255+ self .value = value
256+ self .conversion = conversion
257+ self .expression = expression
258+ self .format_spec = ""
259+
260+
261+ class OurTemplate :
262+ """A simple implementation of IntoTemplate, for testing purposes."""
263+
264+ def __init__ (
265+ self ,
266+ * parts : str | IntoInterpolation ,
267+ ) -> None :
268+ self .strings , self .interpolations = parse_strings_and_params (parts )
269+
270+ def __iter__ (self ) -> Iterator [str | IntoInterpolation ]:
271+ """Iterate over the strings and interpolations in order."""
272+ for s , i in zip (self .strings , self .interpolations , strict = False ):
273+ yield s
274+ yield i
275+ yield self .strings [- 1 ]
276+
277+ def resolve (self ) -> OurTemplate :
278+ """Resolve any inner templates and interpolations, returning a new OurTemplate with only strings and ParamInterpolations."""
279+ return resolve (self )
280+
281+ def compile (self ) -> CompiledSql :
282+ """Compile this template into a final SQL string with named parameter placeholders, and a list of Params."""
283+ resolved = self .resolve ()
284+ return compile_parts (resolved )
285+
286+
287+ T = TypeVar ("T" )
288+
289+
279290def parse_strings_and_params (
280- parts : Iterable [str | IntoParam ],
281- ) -> tuple [tuple [str , ...], tuple [IntoParam , ...]]:
291+ parts : Iterable [str | T ],
292+ ) -> tuple [tuple [str , ...], tuple [T , ...]]:
282293 """Parse an iterable of strings and params into separate tuples of strings and params, merging adjacent strings and ensuring that the number of strings is one more than the number of params."""
283294 strings , params = [], []
284295 last_thing : Literal ["string" , "param" ] | None = None
0 commit comments