From 8dce5065f65f5f24e2780122862ada7deaf2e221 Mon Sep 17 00:00:00 2001
From: Trim21 <trim21.me@gmail.com>
Date: Wed, 22 Nov 2023 13:29:34 +0800
Subject: [PATCH] Add type annotations to xml.py (#1327)

---
 Makefile     |  2 +-
 minio/xml.py | 67 +++++++++++++++++++++++++++++++++++++++-------------
 setup.py     |  3 ++-
 3 files changed, 54 insertions(+), 18 deletions(-)

diff --git a/Makefile b/Makefile
index fc1375653..47638e55a 100644
--- a/Makefile
+++ b/Makefile
@@ -3,7 +3,7 @@ default: tests
 
 getdeps:
 	@echo "Installing required dependencies"
-	@pip install --user --upgrade autopep8 certifi pytest pylint urllib3 argon2-cffi pycryptodome
+	@pip install --user --upgrade autopep8 certifi pytest pylint urllib3 argon2-cffi pycryptodome typing-extensions
 
 check: getdeps
 	@echo "Running checks"
diff --git a/minio/xml.py b/minio/xml.py
index b2727e01d..667f6e00e 100644
--- a/minio/xml.py
+++ b/minio/xml.py
@@ -16,20 +16,28 @@
 
 """XML utility module."""
 
-from __future__ import absolute_import
+from __future__ import absolute_import, annotations
 
 import io
+from typing import Type, TypeVar
 from xml.etree import ElementTree as ET
 
+from typing_extensions import Protocol
+
 _S3_NAMESPACE = "http://s3.amazonaws.com/doc/2006-03-01/"
 
 
-def Element(tag, namespace=_S3_NAMESPACE):  # pylint: disable=invalid-name
+def Element(  # pylint: disable=invalid-name
+    tag: str,
+    namespace: str = _S3_NAMESPACE,
+) -> ET.Element:
     """Create ElementTree.Element with tag and namespace."""
-    return ET.Element(tag, {'xmlns': namespace} if namespace else {})
+    return ET.Element(tag, {"xmlns": namespace} if namespace else {})
 
 
-def SubElement(parent, tag, text=None):  # pylint: disable=invalid-name
+def SubElement(  # pylint: disable=invalid-name
+    parent: ET.Element, tag: str, text: str | None = None
+) -> ET.SubElement:
     """Create ElementTree.SubElement on parent with tag and text."""
     element = ET.SubElement(parent, tag)
     if text is not None:
@@ -37,7 +45,7 @@ def SubElement(parent, tag, text=None):  # pylint: disable=invalid-name
     return element
 
 
-def _get_namespace(element):
+def _get_namespace(element: ET.Element) -> str:
     """Exact namespace if found."""
     start = element.tag.find("{")
     if start < 0:
@@ -49,7 +57,7 @@ def _get_namespace(element):
     return element.tag[start:end]
 
 
-def findall(element, name):
+def findall(element: ET.Element, name: str) -> list[ET.Element]:
     """Namespace aware ElementTree.Element.findall()."""
     namespace = _get_namespace(element)
     return element.findall(
@@ -58,7 +66,7 @@ def findall(element, name):
     )
 
 
-def find(element, name):
+def find(element: ET.Element, name: str) -> ET.Element | None:
     """Namespace aware ElementTree.Element.find()."""
     namespace = _get_namespace(element)
     return element.find(
@@ -67,7 +75,11 @@ def find(element, name):
     )
 
 
-def findtext(element, name, strict=False):
+def findtext(
+    element: ET.Element,
+    name: str,
+    strict: bool = False,
+) -> str | None:
     """
     Namespace aware ElementTree.Element.findtext() with strict flag
     raises ValueError if element name not exist.
@@ -80,20 +92,43 @@ def findtext(element, name, strict=False):
     return element.text or ""
 
 
-def unmarshal(cls, xmlstring):
+K = TypeVar("K")
+
+
+class FromXmlType(Protocol):
+    """typing stub for class with `fromxml` method"""
+
+    @classmethod
+    def fromxml(cls: Type[K], element: ET.Element) -> K:
+        """Create python object with values from XML element."""
+
+
+T = TypeVar("T", bound=FromXmlType)
+
+
+def unmarshal(cls: Type[T], xmlstring: str) -> T:
     """Unmarshal given XML string to an object of passed class."""
     return cls.fromxml(ET.fromstring(xmlstring))
 
 
-def getbytes(element):
+def getbytes(element: ET.Element) -> bytes:
     """Convert ElementTree.Element to bytes."""
-    data = io.BytesIO()
-    ET.ElementTree(element).write(
-        data, encoding=None, xml_declaration=False,
-    )
-    return data.getvalue()
+    with io.BytesIO() as data:
+        ET.ElementTree(element).write(
+            data,
+            encoding=None,
+            xml_declaration=False,
+        )
+        return data.getvalue()
+
+
+class ToXmlType(Protocol):
+    """typing stub for class with `toxml` method"""
+
+    def toxml(self, element: ET.Element | None) -> ET.Element:
+        """Convert python object to ElementTree.Element."""
 
 
-def marshal(obj):
+def marshal(obj: ToXmlType) -> bytes:
     """Get XML data as bytes of ElementTree.Element."""
     return getbytes(obj.toxml(None))
diff --git a/setup.py b/setup.py
index 28b01e8ea..4a888ffef 100644
--- a/setup.py
+++ b/setup.py
@@ -43,7 +43,8 @@
     long_description_content_type="text/markdown",
     package_dir={"minio": "minio"},
     packages=["minio", "minio.credentials"],
-    install_requires=["certifi", "urllib3", "argon2-cffi", "pycryptodome"],
+    install_requires=["certifi", "urllib3", "argon2-cffi",
+                      "pycryptodome", "typing-extensions"],
     tests_require=[],
     license="Apache-2.0",
     classifiers=[