codekingpro/portable-devtools

1import io
2import warnings
3from typing import Any, Iterable, List, Optional
4from urllib.parse import urlencode
5
6from multidict import MultiDict, MultiDictProxy
7
8from . import hdrs, multipart, payload
9from .helpers import guess_filename
10from .payload import Payload
11
12__all__ = ("FormData",)
13
14
15class FormData:
16    """Helper class for form body generation.
17
18    Supports multipart/form-data and application/x-www-form-urlencoded.
19    """
20
21    def __init__(
22        self,
23        fields: Iterable[Any] = (),
24        quote_fields: bool = True,
25        charset: Optional[str] = None,
26        *,
27        default_to_multipart: bool = False,
28    ) -> None:
29        self._writer = multipart.MultipartWriter("form-data")
30        self._fields: List[Any] = []
31        self._is_multipart = default_to_multipart
32        self._quote_fields = quote_fields
33        self._charset = charset
34
35        if isinstance(fields, dict):
36            fields = list(fields.items())
37        elif not isinstance(fields, (list, tuple)):
38            fields = (fields,)
39        self.add_fields(*fields)
40
41    @property
42    def is_multipart(self) -> bool:
43        return self._is_multipart
44
45    def add_field(
46        self,
47        name: str,
48        value: Any,
49        *,
50        content_type: Optional[str] = None,
51        filename: Optional[str] = None,
52        content_transfer_encoding: Optional[str] = None,
53    ) -> None:
54
55        if isinstance(value, io.IOBase):
56            self._is_multipart = True
57        elif isinstance(value, (bytes, bytearray, memoryview)):
58            msg = (
59                "In v4, passing bytes will no longer create a file field. "
60                "Please explicitly use the filename parameter or pass a BytesIO object."
61            )
62            if filename is None and content_transfer_encoding is None:
63                warnings.warn(msg, DeprecationWarning)
64                filename = name
65
66        type_options: MultiDict[str] = MultiDict({"name": name})
67        if filename is not None and not isinstance(filename, str):
68            raise TypeError("filename must be an instance of str. Got: %s" % filename)
69        if filename is None and isinstance(value, io.IOBase):
70            filename = guess_filename(value, name)
71        if filename is not None:
72            type_options["filename"] = filename
73            self._is_multipart = True
74
75        headers = {}
76        if content_type is not None:
77            if not isinstance(content_type, str):
78                raise TypeError(
79                    "content_type must be an instance of str. Got: %s" % content_type
80                )
81            if "\r" in content_type or "\n" in content_type:
82                raise ValueError(
83                    "Newline or carriage return detected in headers. "
84                    "Potential header injection attack."
85                )
86            headers[hdrs.CONTENT_TYPE] = content_type
87            self._is_multipart = True
88        if content_transfer_encoding is not None:
89            if not isinstance(content_transfer_encoding, str):
90                raise TypeError(
91                    "content_transfer_encoding must be an instance"
92                    " of str. Got: %s" % content_transfer_encoding
93                )
94            msg = (
95                "content_transfer_encoding is deprecated. "
96                "To maintain compatibility with v4 please pass a BytesPayload."
97            )
98            warnings.warn(msg, DeprecationWarning)
99            self._is_multipart = True
100
101        self._fields.append((type_options, headers, value))
102
103    def add_fields(self, *fields: Any) -> None:
104        to_add = list(fields)
105
106        while to_add:
107            rec = to_add.pop(0)
108
109            if isinstance(rec, io.IOBase):
110                k = guess_filename(rec, "unknown")
111                self.add_field(k, rec)  # type: ignore[arg-type]
112
113            elif isinstance(rec, (MultiDictProxy, MultiDict)):
114                to_add.extend(rec.items())
115
116            elif isinstance(rec, (list, tuple)) and len(rec) == 2:
117                k, fp = rec
118                self.add_field(k, fp)
119
120            else:
121                raise TypeError(
122                    "Only io.IOBase, multidict and (name, file) "
123                    "pairs allowed, use .add_field() for passing "
124                    "more complex parameters, got {!r}".format(rec)
125                )
126
127    def _gen_form_urlencoded(self) -> payload.BytesPayload:
128        # form data (x-www-form-urlencoded)
129        data = []
130        for type_options, _, value in self._fields:
131            data.append((type_options["name"], value))
132
133        charset = self._charset if self._charset is not None else "utf-8"
134
135        if charset == "utf-8":
136            content_type = "application/x-www-form-urlencoded"
137        else:
138            content_type = "application/x-www-form-urlencoded; charset=%s" % charset
139
140        return payload.BytesPayload(
141            urlencode(data, doseq=True, encoding=charset).encode(),
142            content_type=content_type,
143        )
144
145    def _gen_form_data(self) -> multipart.MultipartWriter:
146        """Encode a list of fields using the multipart/form-data MIME format"""
147        for dispparams, headers, value in self._fields:
148            try:
149                if hdrs.CONTENT_TYPE in headers:
150                    part = payload.get_payload(
151                        value,
152                        content_type=headers[hdrs.CONTENT_TYPE],
153                        headers=headers,
154                        encoding=self._charset,
155                    )
156                else:
157                    part = payload.get_payload(
158                        value, headers=headers, encoding=self._charset
159                    )
160            except Exception as exc:
161                raise TypeError(
162                    "Can not serialize value type: %r\n "
163                    "headers: %r\n value: %r" % (type(value), headers, value)
164                ) from exc
165
166            if dispparams:
167                part.set_content_disposition(
168                    "form-data", quote_fields=self._quote_fields, **dispparams
169                )
170                # FIXME cgi.FieldStorage doesn't likes body parts with
171                # Content-Length which were sent via chunked transfer encoding
172                assert part.headers is not None
173                part.headers.popall(hdrs.CONTENT_LENGTH, None)
174
175            self._writer.append_payload(part)
176
177        self._fields.clear()
178        return self._writer
179
180    def __call__(self) -> Payload:
181        if self._is_multipart:
182            return self._gen_form_data()
183        else:
184            return self._gen_form_urlencoded()
185