Support Python Literal types (#78)

.pre-commit-config.yaml

@@ -15,28 +15,28 @@
 
 repos:
   - repo: https://github.com/psf/black
-    rev: 23.3.0
+    rev: 24.8.0
     hooks:
       - # Format code
         id: black
   - repo: https://github.com/PyCQA/flake8
-    rev: 6.0.0
+    rev: 7.1.1
     hooks:
       - # Apply linting rules
         id: flake8
   - repo: https://github.com/econchick/interrogate
-    rev: 1.5.0
+    rev: 1.7.0
     hooks:
       - # Enforce documentation
         id: interrogate
         exclude: ^tests
   - repo: https://github.com/PyCQA/isort
-    rev: 5.12.0
+    rev: 5.13.2
     hooks:
       - # Sort imports
         id: isort
   - repo: https://github.com/Lucas-C/pre-commit-hooks
-    rev: v1.5.1
+    rev: v1.5.5
     hooks:
       - # Insert an OSS license header to all relevant files
         id: insert-license

src/py_avro_schema/_schemas.py

@@ -34,7 +34,9 @@
     Dict,
     ForwardRef,
     List,
+    Literal,
     Optional,
+    Set,
     Tuple,
     Type,
     Union,
@@ -309,6 +311,41 @@ def data(self, names: NamesType) -> JSONObj:
         }
 
 
+class LiteralSchema(Schema):
+    """An Avro schema of any type for a Python Literal type, e.g. Literal[""]"""
+
+    def __init__(self, py_type: Type[Any], namespace: Optional[str] = None, options: Option = Option(0)):
+        """
+        An Avro schema of any type for a Python Literal type, e.g. Literal[""]
+
+        :param py_type:   The Python class to generate a schema for.
+        :param namespace: The Avro namespace to add to schemas.
+        :param options:   Schema generation options.
+        """
+        super().__init__(py_type, namespace=namespace, options=options)
+        py_type = _type_from_annotated(py_type)
+        literal_type = self._literal_value_types(py_type).pop()
+        self.literal_value_schema = _schema_obj(literal_type, namespace=namespace, options=options)
+
+    @classmethod
+    def handles_type(cls, py_type: Type[Any]) -> bool:
+        """Whether this schema class can represent a given Python class"""
+        py_type = _type_from_annotated(py_type)
+        literal_value_types = cls._literal_value_types(py_type)
+        # For now we support Literals with the same type only. Potentially we could explose multiple literal types into
+        # an Avro Union schema, but that may not be something that anyone would every want to use...
+        return get_origin(py_type) is Literal and len(literal_value_types) == 1
+
+    def data(self, names: NamesType) -> JSONType:
+        """Return the schema data"""
+        return self.literal_value_schema.data(names=names)
+
+    @staticmethod
+    def _literal_value_types(py_type) -> Set[Type[Any]]:
+        """Return the Python types corresponding to the literal values"""
+        return {type(literal_value) for literal_value in get_args(py_type)}
+
+
 class DictAsJSONSchema(Schema):
     """An Avro string schema representing a Python Dict[str, Any] or List[Dict[str, Any]] assuming JSON serialization"""
 

tests/test_primitives.py

@@ -16,6 +16,7 @@
     Annotated,
     Dict,
     List,
+    Literal,
     Mapping,
     MutableSequence,
     Optional,
@@ -75,6 +76,24 @@ class PyType(packaging.version.Version, str): ...
     assert_schema(PyType, expected)
 
 
+def test_str_literal():
+    py_type = Literal[""]
+    expected = "string"
+    assert_schema(py_type, expected)
+
+
+def test_str_literal_multiple():
+    py_type = Literal["", "Hello, Python"]
+    expected = "string"
+    assert_schema(py_type, expected)
+
+
+def test_str_literal_annotated():
+    py_type = Annotated[Literal[""], ...]
+    expected = "string"
+    assert_schema(py_type, expected)
+
+
 def test_int():
     py_type = int
     expected = "long"
@@ -94,6 +113,12 @@ def test_int_32():
     assert_schema(py_type, expected, options=options)
 
 
+def test_int_literal():
+    py_type = Literal[42]
+    expected = "long"
+    assert_schema(py_type, expected)
+
+
 def test_bool():
     py_type = bool
     expected = "boolean"
@@ -327,6 +352,13 @@ def test_union_of_union_string_int():
     assert_schema(py_type, expected)
 
 
+def test_literal_different_types():
+    py_type = Literal["", 42]
+    expected = {}
+    with pytest.raises(pas.TypeNotSupportedError):
+        assert_schema(py_type, expected)
+
+
 def test_optional_str():
     py_type = Optional[str]
     expected = ["string", "null"]