import base64
import binascii
from typing import Any
from pydantic import GetCoreSchemaHandler
from pydantic_core import CoreSchema, core_schema
[docs]
class Base64Data:
[docs]
def __init__(self, data: bytes):
"""
Initializes the object.
If the provided data is already in base64 encoded format, it will store it.
If the data is a regular byte string, it will encode and then store it.
"""
if self.is_base64(data):
self._data = str(data, "utf-8")
else:
encoded = base64.b64encode(data)
self._data = str(encoded, "utf-8")
[docs]
@staticmethod
def is_base64(data: bytes) -> bool:
"""Checks if given data is base64 encoded."""
try:
str_data = str(data, "utf-8")
_ = base64.urlsafe_b64decode(str_data.strip("=") + "===")
return True
except binascii.Error:
return False
[docs]
def get_encoded(self) -> str:
"""Returns the stored base64 encoded data."""
return self._data
[docs]
def get_decoded(self) -> bytes:
"""Returns the decoded byte string."""
return base64.urlsafe_b64decode(self._data.strip("=") + "===")
[docs]
@classmethod
def __get_pydantic_core_schema__(
cls, source_type: Any, handler: GetCoreSchemaHandler
) -> CoreSchema:
return core_schema.no_info_after_validator_function(cls, handler(bytes))