codekingpro/portable-devtools

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