codec.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159
  1. import codecs
  2. from typing import Any, Optional
  3. from .core import IDNAError, _unicode_dots_re, alabel, decode, encode, ulabel
  4. class Codec(codecs.Codec):
  5. """Stateless IDNA 2008 codec.
  6. Implements the :class:`codecs.Codec` protocol so that the whole-domain
  7. encoder (:func:`idna.encode`) and decoder (:func:`idna.decode`) are
  8. accessible through the standard codec machinery as ``"idna2008"``.
  9. Only the ``"strict"`` error handler is supported; any other handler
  10. raises :exc:`~idna.IDNAError`.
  11. """
  12. def encode(self, data: str, errors: str = "strict") -> tuple[bytes, int]: # ty: ignore[invalid-method-override]
  13. if errors != "strict":
  14. raise IDNAError(f'Unsupported error handling "{errors}"')
  15. if not data:
  16. return b"", 0
  17. return encode(data), len(data)
  18. def decode(self, data: bytes, errors: str = "strict") -> tuple[str, int]: # ty: ignore[invalid-method-override]
  19. if errors != "strict":
  20. raise IDNAError(f'Unsupported error handling "{errors}"')
  21. if not data:
  22. return "", 0
  23. return decode(data), len(data)
  24. class IncrementalEncoder(codecs.BufferedIncrementalEncoder):
  25. """Incremental IDNA 2008 encoder.
  26. Buffers a partial trailing label across calls until either the next
  27. label separator is seen or ``final=True``, so that streamed input is
  28. encoded one whole label at a time. Any of the four Unicode label
  29. separators (``U+002E``, ``U+3002``, ``U+FF0E``, ``U+FF61``) ends a
  30. label; the result always uses ``U+002E`` as the separator.
  31. Only the ``"strict"`` error handler is supported.
  32. """
  33. def _buffer_encode(self, data: str, errors: str, final: bool) -> tuple[bytes, int]: # ty: ignore[invalid-method-override]
  34. if errors != "strict":
  35. raise IDNAError(f'Unsupported error handling "{errors}"')
  36. if not data:
  37. return b"", 0
  38. labels = _unicode_dots_re.split(data)
  39. trailing_dot = b""
  40. if labels:
  41. if not labels[-1]:
  42. trailing_dot = b"."
  43. del labels[-1]
  44. elif not final:
  45. # Keep potentially unfinished label until the next call
  46. del labels[-1]
  47. if labels:
  48. trailing_dot = b"."
  49. result = []
  50. size = 0
  51. for label in labels:
  52. result.append(alabel(label))
  53. if size:
  54. size += 1
  55. size += len(label)
  56. # Join with U+002E
  57. result_bytes = b".".join(result) + trailing_dot
  58. size += len(trailing_dot)
  59. return result_bytes, size
  60. class IncrementalDecoder(codecs.BufferedIncrementalDecoder):
  61. """Incremental IDNA 2008 decoder.
  62. Buffers a partial trailing label across calls until either the next
  63. label separator is seen or ``final=True``, so that streamed input is
  64. decoded one whole label at a time.
  65. Only the ``"strict"`` error handler is supported.
  66. """
  67. def _buffer_decode(self, data: Any, errors: str, final: bool) -> tuple[str, int]: # ty: ignore[invalid-method-override]
  68. if errors != "strict":
  69. raise IDNAError(f'Unsupported error handling "{errors}"')
  70. if not data:
  71. return ("", 0)
  72. if not isinstance(data, str):
  73. data = str(data, "ascii")
  74. labels = _unicode_dots_re.split(data)
  75. trailing_dot = ""
  76. if labels:
  77. if not labels[-1]:
  78. trailing_dot = "."
  79. del labels[-1]
  80. elif not final:
  81. # Keep potentially unfinished label until the next call
  82. del labels[-1]
  83. if labels:
  84. trailing_dot = "."
  85. result = []
  86. size = 0
  87. for label in labels:
  88. result.append(ulabel(label))
  89. if size:
  90. size += 1
  91. size += len(label)
  92. result_str = ".".join(result) + trailing_dot
  93. size += len(trailing_dot)
  94. return (result_str, size)
  95. class StreamWriter(Codec, codecs.StreamWriter):
  96. pass
  97. class StreamReader(Codec, codecs.StreamReader):
  98. pass
  99. def search_function(name: str) -> Optional[codecs.CodecInfo]:
  100. """Codec search function registered with :mod:`codecs`.
  101. Returns a :class:`codecs.CodecInfo` for the ``"idna2008"`` codec name
  102. so that ``str.encode("idna2008")`` and ``bytes.decode("idna2008")``
  103. invoke the IDNA 2008 codec defined in this module.
  104. :param name: The codec name being looked up.
  105. :returns: A :class:`codecs.CodecInfo` instance if ``name`` is
  106. ``"idna2008"``, otherwise ``None``.
  107. """
  108. if name != "idna2008":
  109. return None
  110. return codecs.CodecInfo(
  111. name=name,
  112. encode=Codec().encode,
  113. decode=Codec().decode, # type: ignore
  114. incrementalencoder=IncrementalEncoder,
  115. incrementaldecoder=IncrementalDecoder,
  116. streamwriter=StreamWriter,
  117. streamreader=StreamReader,
  118. )
  119. codecs.register(search_function)