""" 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}_]" )