codekingpro/portable-devtools

1# dialects/mysql/enumerated.py
2# Copyright (C) 2005-2026 the SQLAlchemy authors and contributors
3# <see AUTHORS file>
4#
5# This module is part of SQLAlchemy and is released under
6# the MIT License: https://www.opensource.org/licenses/mit-license.php
7
8from __future__ import annotations
9
10import enum
11import re
12from typing import Any
13from typing import Dict
14from typing import Optional
15from typing import Set
16from typing import Type
17from typing import TYPE_CHECKING
18from typing import Union
19
20from .types import _StringType
21from ... import exc
22from ... import sql
23from ... import util
24from ...sql import sqltypes
25from ...sql import type_api
26
27if TYPE_CHECKING:
28    from ...engine.interfaces import Dialect
29    from ...sql.elements import ColumnElement
30    from ...sql.type_api import _BindProcessorType
31    from ...sql.type_api import _ResultProcessorType
32    from ...sql.type_api import TypeEngine
33    from ...sql.type_api import TypeEngineMixin
34
35
36class ENUM(type_api.NativeForEmulated, sqltypes.Enum, _StringType):
37    """MySQL ENUM type."""
38
39    __visit_name__ = "ENUM"
40
41    native_enum = True
42
43    def __init__(self, *enums: Union[str, Type[enum.Enum]], **kw: Any) -> None:
44        """Construct an ENUM.
45
46        E.g.::
47
48          Column("myenum", ENUM("foo", "bar", "baz"))
49
50        :param enums: The range of valid values for this ENUM.  Values in
51          enums are not quoted, they will be escaped and surrounded by single
52          quotes when generating the schema.  This object may also be a
53          PEP-435-compliant enumerated type.
54
55          .. versionadded: 1.1 added support for PEP-435-compliant enumerated
56             types.
57
58        :param strict: This flag has no effect.
59
60         .. versionchanged:: The MySQL ENUM type as well as the base Enum
61            type now validates all Python data values.
62
63        :param charset: Optional, a column-level character set for this string
64          value.  Takes precedence to 'ascii' or 'unicode' short-hand.
65
66        :param collation: Optional, a column-level collation for this string
67          value.  Takes precedence to 'binary' short-hand.
68
69        :param ascii: Defaults to False: short-hand for the ``latin1``
70          character set, generates ASCII in schema.
71
72        :param unicode: Defaults to False: short-hand for the ``ucs2``
73          character set, generates UNICODE in schema.
74
75        :param binary: Defaults to False: short-hand, pick the binary
76          collation type that matches the column's character set.  Generates
77          BINARY in schema.  This does not affect the type of data stored,
78          only the collation of character data.
79
80        """
81        kw.pop("strict", None)
82        self._enum_init(enums, kw)  # type: ignore[arg-type]
83        _StringType.__init__(self, length=self.length, **kw)
84
85    @classmethod
86    def adapt_emulated_to_native(
87        cls,
88        impl: Union[TypeEngine[Any], TypeEngineMixin],
89        **kw: Any,
90    ) -> ENUM:
91        """Produce a MySQL native :class:`.mysql.ENUM` from plain
92        :class:`.Enum`.
93
94        """
95        if TYPE_CHECKING:
96            assert isinstance(impl, ENUM)
97        kw.setdefault("validate_strings", impl.validate_strings)
98        kw.setdefault("values_callable", impl.values_callable)
99        kw.setdefault("omit_aliases", impl._omit_aliases)
100        return cls(**kw)
101
102    def _object_value_for_elem(self, elem: str) -> Union[str, enum.Enum]:
103        # mysql sends back a blank string for any value that
104        # was persisted that was not in the enums; that is, it does no
105        # validation on the incoming data, it "truncates" it to be
106        # the blank string.  Return it straight.
107        if elem == "":
108            return elem
109        else:
110            return super()._object_value_for_elem(elem)
111
112    def __repr__(self) -> str:
113        return util.generic_repr(
114            self, to_inspect=[ENUM, _StringType, sqltypes.Enum]
115        )
116
117
118# TODO: SET is a string as far as configuration but does not act like
119# a string at the python level.  We either need to make a py-type agnostic
120# version of String as a base to be used for this, make this some kind of
121# TypeDecorator, or just vendor it out as its own type.
122class SET(_StringType):
123    """MySQL SET type."""
124
125    __visit_name__ = "SET"
126
127    def __init__(self, *values: str, **kw: Any):
128        """Construct a SET.
129
130        E.g.::
131
132          Column("myset", SET("foo", "bar", "baz"))
133
134        The list of potential values is required in the case that this
135        set will be used to generate DDL for a table, or if the
136        :paramref:`.SET.retrieve_as_bitwise` flag is set to True.
137
138        :param values: The range of valid values for this SET. The values
139          are not quoted, they will be escaped and surrounded by single
140          quotes when generating the schema.
141
142        :param convert_unicode: Same flag as that of
143         :paramref:`.String.convert_unicode`.
144
145        :param collation: same as that of :paramref:`.String.collation`
146
147        :param charset: same as that of :paramref:`.VARCHAR.charset`.
148
149        :param ascii: same as that of :paramref:`.VARCHAR.ascii`.
150
151        :param unicode: same as that of :paramref:`.VARCHAR.unicode`.
152
153        :param binary: same as that of :paramref:`.VARCHAR.binary`.
154
155        :param retrieve_as_bitwise: if True, the data for the set type will be
156          persisted and selected using an integer value, where a set is coerced
157          into a bitwise mask for persistence.  MySQL allows this mode which
158          has the advantage of being able to store values unambiguously,
159          such as the blank string ``''``.   The datatype will appear
160          as the expression ``col + 0`` in a SELECT statement, so that the
161          value is coerced into an integer value in result sets.
162          This flag is required if one wishes
163          to persist a set that can store the blank string ``''`` as a value.
164
165          .. warning::
166
167            When using :paramref:`.mysql.SET.retrieve_as_bitwise`, it is
168            essential that the list of set values is expressed in the
169            **exact same order** as exists on the MySQL database.
170
171        """
172        self.retrieve_as_bitwise = kw.pop("retrieve_as_bitwise", False)
173        self.values = tuple(values)
174        if not self.retrieve_as_bitwise and "" in values:
175            raise exc.ArgumentError(
176                "Can't use the blank value '' in a SET without "
177                "setting retrieve_as_bitwise=True"
178            )
179        if self.retrieve_as_bitwise:
180            self._inversed_bitmap: Dict[str, int] = {
181                value: 2**idx for idx, value in enumerate(self.values)
182            }
183            self._bitmap: Dict[int, str] = {
184                2**idx: value for idx, value in enumerate(self.values)
185            }
186        length = max([len(v) for v in values] + [0])
187        kw.setdefault("length", length)
188        super().__init__(**kw)
189
190    def column_expression(
191        self, colexpr: ColumnElement[Any]
192    ) -> ColumnElement[Any]:
193        if self.retrieve_as_bitwise:
194            return sql.type_coerce(
195                sql.type_coerce(colexpr, sqltypes.Integer) + 0, self
196            )
197        else:
198            return colexpr
199
200    def result_processor(
201        self, dialect: Dialect, coltype: Any
202    ) -> Optional[_ResultProcessorType[Any]]:
203        if self.retrieve_as_bitwise:
204
205            def process(value: Union[str, int, None]) -> Optional[Set[str]]:
206                if value is not None:
207                    value = int(value)
208
209                    return set(util.map_bits(self._bitmap.__getitem__, value))
210                else:
211                    return None
212
213        else:
214            super_convert = super().result_processor(dialect, coltype)
215
216            def process(value: Union[str, Set[str], None]) -> Optional[Set[str]]:  # type: ignore[misc]  # noqa: E501
217                if isinstance(value, str):
218                    # MySQLdb returns a string, let's parse
219                    if super_convert:
220                        value = super_convert(value)
221                        assert value is not None
222                    if TYPE_CHECKING:
223                        assert isinstance(value, str)
224                    return set(re.findall(r"[^,]+", value))
225                else:
226                    # mysql-connector-python does a naive
227                    # split(",") which throws in an empty string
228                    if value is not None:
229                        value.discard("")
230                    return value
231
232        return process
233
234    def bind_processor(
235        self, dialect: Dialect
236    ) -> _BindProcessorType[Union[str, int]]:
237        super_convert = super().bind_processor(dialect)
238        if self.retrieve_as_bitwise:
239
240            def process(
241                value: Union[str, int, set[str], None],
242            ) -> Union[str, int, None]:
243                if value is None:
244                    return None
245                elif isinstance(value, (int, str)):
246                    if super_convert:
247                        return super_convert(value)  # type: ignore[arg-type, no-any-return]  # noqa: E501
248                    else:
249                        return value
250                else:
251                    int_value = 0
252                    for v in value:
253                        int_value |= self._inversed_bitmap[v]
254                    return int_value
255
256        else:
257
258            def process(
259                value: Union[str, int, set[str], None],
260            ) -> Union[str, int, None]:
261                # accept strings and int (actually bitflag) values directly
262                if value is not None and not isinstance(value, (int, str)):
263                    value = ",".join(value)
264                if super_convert:
265                    return super_convert(value)  # type: ignore
266                else:
267                    return value
268
269        return process
270
271    def adapt(self, cls: type, **kw: Any) -> Any:
272        kw["retrieve_as_bitwise"] = self.retrieve_as_bitwise
273        return util.constructor_copy(self, cls, *self.values, **kw)
274
275    def __repr__(self) -> str:
276        return util.generic_repr(
277            self,
278            to_inspect=[SET, _StringType],
279            additional_kw=[
280                ("retrieve_as_bitwise", False),
281            ],
282        )
283