153 lines
4.5 KiB
Python
153 lines
4.5 KiB
Python
"""
|
|
Mappings between PostgreSQL and Python encodings.
|
|
"""
|
|
|
|
# Copyright (C) 2020 The Psycopg Team
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
import codecs
|
|
import string
|
|
from typing import TYPE_CHECKING, Any
|
|
from functools import cache
|
|
|
|
from .errors import NotSupportedError
|
|
|
|
if TYPE_CHECKING:
|
|
from ._connection_base import BaseConnection
|
|
|
|
# From: https://www.postgresql.org/docs/current/multibyte.html#CHARSET-TABLE
|
|
_py_codecs = (
|
|
(("BIG5", "WIN950", "Windows950"), "big5"),
|
|
(("EUC_CN",), "gb2312"),
|
|
(("EUC_JIS_2004",), "euc_jis_2004"),
|
|
(("EUC_JP",), "euc_jp"),
|
|
(("EUC_KR",), "euc_kr"),
|
|
# (("EUC_TW",), not available in Python
|
|
(("GB18030",), "gb18030"),
|
|
(("GBK", "WIN936", "Windows936"), "gbk"),
|
|
(("ISO_8859_5",), "iso8859-5"),
|
|
(("ISO_8859_6",), "iso8859-6"),
|
|
(("ISO_8859_7",), "iso8859-7"),
|
|
(("ISO_8859_8",), "iso8859-8"),
|
|
(("JOHAB",), "johab"),
|
|
(("KOI8R", "KOI8"), "koi8-r"),
|
|
(("KOI8U",), "koi8-u"),
|
|
(("LATIN1", "ISO88591"), "iso8859-1"),
|
|
(("LATIN10", "ISO885916"), "iso8859-16"),
|
|
(("LATIN2", "ISO88592"), "iso8859-2"),
|
|
(("LATIN3", "ISO88593"), "iso8859-3"),
|
|
(("LATIN4", "ISO88594"), "iso8859-4"),
|
|
(("LATIN5", "ISO88599"), "iso8859-9"),
|
|
(("LATIN6", "ISO885910"), "iso8859-10"),
|
|
(("LATIN7", "ISO885913"), "iso8859-13"),
|
|
(("LATIN8", "ISO885914"), "iso8859-14"),
|
|
(("LATIN9", "ISO885915"), "iso8859-15"),
|
|
# "MULE_INTERNAL" not available in Python
|
|
(("SHIFT_JIS_2004",), "shift_jis_2004"),
|
|
(("SJIS", "Mskanji", "ShiftJIS", "WIN932", "Windows932"), "shift_jis"),
|
|
# this actually means no encoding, see PostgreSQL docs
|
|
# it is special-cased by the text loader.
|
|
(("SQL_ASCII",), "ascii"),
|
|
(("UHC", "WIN949", "Windows949"), "cp949"),
|
|
(("UTF8", "Unicode"), "utf-8"),
|
|
(("WIN1250",), "cp1250"),
|
|
(("WIN1251", "WIN"), "cp1251"),
|
|
(("WIN1252",), "cp1252"),
|
|
(("WIN1253",), "cp1253"),
|
|
(("WIN1254",), "cp1254"),
|
|
(("WIN1255",), "cp1255"),
|
|
(("WIN1256",), "cp1256"),
|
|
(("WIN1257",), "cp1257"),
|
|
(("WIN1258", "ABC", "TCVN", "TCVN5712", "VSCII"), "cp1258"),
|
|
(("WIN866", "ALT"), "cp866"),
|
|
(("WIN874",), "cp874"),
|
|
)
|
|
|
|
py_codecs = {a.upper().encode(): v for aliases, v in _py_codecs for a in aliases}
|
|
|
|
# Add an alias without underscore, for lenient lookups
|
|
py_codecs.update(
|
|
(a.replace("_", "").encode(), v)
|
|
for aliases, v in _py_codecs
|
|
for a in aliases
|
|
if "_" in a
|
|
)
|
|
|
|
pg_codecs = {v: k[0].encode() for k, v in _py_codecs}
|
|
|
|
|
|
def conn_encoding(conn: BaseConnection[Any] | None) -> str:
|
|
"""
|
|
Return the Python encoding name of a psycopg connection.
|
|
|
|
Default to utf8 if the connection has no encoding info.
|
|
"""
|
|
return conn.pgconn._encoding if conn else "utf-8"
|
|
|
|
|
|
def conninfo_encoding(conninfo: str) -> str:
|
|
"""
|
|
Return the Python encoding name passed in a conninfo string. Default to utf8.
|
|
|
|
Because the input is likely to come from the user and not normalised by the
|
|
server, be somewhat lenient (non-case-sensitive lookup, ignore noise chars).
|
|
"""
|
|
from .conninfo import conninfo_to_dict
|
|
|
|
params = conninfo_to_dict(conninfo)
|
|
if pgenc := params.get("client_encoding"):
|
|
try:
|
|
return pg2pyenc(str(pgenc).encode())
|
|
except NotSupportedError:
|
|
pass
|
|
|
|
return "utf-8"
|
|
|
|
|
|
@cache
|
|
def py2pgenc(name: str) -> bytes:
|
|
"""Convert a Python encoding name to PostgreSQL encoding name.
|
|
|
|
Raise LookupError if the Python encoding is unknown.
|
|
"""
|
|
return pg_codecs[codecs.lookup(name).name]
|
|
|
|
|
|
@cache
|
|
def pg2pyenc(name: bytes) -> str:
|
|
"""Convert a PostgreSQL encoding name to Python encoding name.
|
|
|
|
Raise NotSupportedError if the PostgreSQL encoding is not supported by
|
|
Python.
|
|
"""
|
|
try:
|
|
return py_codecs[name.replace(b"-", b"").replace(b"_", b"").upper()]
|
|
except KeyError:
|
|
sname = name.decode("utf8", "replace")
|
|
raise NotSupportedError(f"codec not available in Python: {sname!r}")
|
|
|
|
|
|
def _as_python_identifier(s: str, prefix: str = "f") -> str:
|
|
"""
|
|
Reduce a string to a valid Python identifier.
|
|
|
|
Replace all non-valid chars with '_' and prefix the value with `!prefix` if
|
|
the first letter is an '_'.
|
|
"""
|
|
if not s.isidentifier():
|
|
if s[0] in "1234567890":
|
|
s = prefix + s
|
|
if not s.isidentifier():
|
|
s = _re_clean.sub("_", s)
|
|
# namedtuple fields cannot start with underscore. So...
|
|
if s[0] == "_":
|
|
s = prefix + s
|
|
return s
|
|
|
|
|
|
_re_clean = re.compile(
|
|
f"[^{string.ascii_lowercase}{string.ascii_uppercase}{string.digits}_]"
|
|
)
|