@@ -1595,20 +1595,113 @@ def MkTempFile(data=None, inmem=__use_inmemfile__, isbytes=True, prefix=__projec
15951595 return f
15961596
15971597
1598- def MkTempFileSmart (data=None , isbytes=True, prefix=__project__, max_mem=1024*1024, encoding="utf-8" ):
1598+ def _normalize_initial_data (data, isbytes, encoding):
15991599 """
1600- Spooled temp file: starts in memory and spills to disk past max_mem.
1601- Behaves like BytesIO/StringIO for small data, with the same preload+seek(0) behavior.
1600+ Coerce `data` to the correct type for the chosen mode:
1601+ - bytes mode: return `bytes` (Py2: str; Py3: bytes)
1602+ - text mode : return unicode/str (Py2: unicode; Py3: str)
16021603 """
1603- mode = "wb+" if isbytes else "w+"
1604- kwargs = {"mode": mode, "max_size": max_mem, "prefix": prefix or ""}
1605- if not isbytes and sys.version_info[0] >= 3:
1606- kwargs["encoding"] = encoding
1607- kwargs["newline"] = ""
1604+ if data is None:
1605+ return None
16081606
1609- f = tempfile.SpooledTemporaryFile(**kwargs)
1607+ if isbytes:
1608+ # Need a byte sequence
1609+ if isinstance(data, bytes):
1610+ return data
1611+ if isinstance(data, bytearray):
1612+ return bytes(data)
1613+ # memoryview may not exist on very old Py2 builds; guard dynamically
1614+ mv_t = getattr(__builtins__, 'memoryview', type(None))
1615+ if isinstance(data, mv_t):
1616+ return bytes(data)
1617+ if isinstance(data, str):
1618+ # Py2 str is already bytes; Py3 str must be encoded
1619+ return data if PY2 else data.encode(encoding)
1620+ if PY2 and isinstance(data, unicode): # noqa: F821 (unicode only in Py2)
1621+ return data.encode(encoding)
1622+ raise TypeError("data must be bytes-like or text for isbytes=True (got %r)" % (type(data),))
1623+ else:
1624+ # Need text (unicode in Py2, str in Py3)
1625+ if PY2:
1626+ if isinstance(data, unicode): # noqa: F821
1627+ return data
1628+ if isinstance(data, str):
1629+ return data.decode(encoding)
1630+ if isinstance(data, bytearray):
1631+ return bytes(data).decode(encoding)
1632+ mv_t = getattr(__builtins__, 'memoryview', type(None))
1633+ if isinstance(data, mv_t):
1634+ return bytes(data).decode(encoding)
1635+ raise TypeError("data must be unicode or bytes-like for text mode (got %r)" % (type(data),))
1636+ else:
1637+ if isinstance(data, str):
1638+ return data
1639+ if isinstance(data, (bytes, bytearray, memoryview)):
1640+ return bytes(data).decode(encoding)
1641+ raise TypeError("data must be str or bytes-like for text mode (got %r)" % (type(data),))
1642+
1643+ def MkTempFile(data=None,
1644+ inmem=True,
1645+ isbytes=True,
1646+ prefix="",
1647+ delete=True,
1648+ encoding="utf-8",
1649+ newline=None, # Py3 text only; ignored by Py2 temp classes
1650+ dir=None,
1651+ suffix="",
1652+ # spooled option (RAM until threshold, then rolls to disk)
1653+ use_spool=False,
1654+ spool_max=8 * 1024 * 1024,
1655+ spool_dir=None):
1656+ """
1657+ Return a file-like handle with consistent behavior on Py2.7 and Py3.x.
16101658
1659+ - inmem=True -> BytesIO (bytes) or StringIO (text)
1660+ - inmem=False, use_spool=True -> SpooledTemporaryFile (RAM -> disk after spool_max)
1661+ - inmem=False, use_spool=False -> NamedTemporaryFile (on disk)
1662+
1663+ If `data` is provided, it's written and the handle is rewound to position 0.
1664+ """
16111665 init = _normalize_initial_data(data, isbytes, encoding)
1666+
1667+ # -------- In-memory --------
1668+ if inmem:
1669+ if isbytes:
1670+ return BytesIO(init if init is not None else b"")
1671+ else:
1672+ # Py2 needs unicode literal for empty default
1673+ return StringIO(init if init is not None else (u"" if PY2 else ""))
1674+
1675+ # -------- Spooled (RAM then disk) --------
1676+ if use_spool:
1677+ if isbytes:
1678+ f = tempfile.SpooledTemporaryFile(max_size=spool_max, mode="w+b", dir=spool_dir)
1679+ else:
1680+ if PY2:
1681+ # Py2 SpooledTemporaryFile doesn't accept encoding/newline
1682+ f = tempfile.SpooledTemporaryFile(max_size=spool_max, mode="w+", dir=spool_dir)
1683+ else:
1684+ f = tempfile.SpooledTemporaryFile(max_size=spool_max, mode="w+",
1685+ dir=spool_dir, encoding=encoding, newline=newline)
1686+ if init is not None:
1687+ f.write(init)
1688+ f.seek(0)
1689+ return f
1690+
1691+ # -------- On-disk temp --------
1692+ if isbytes:
1693+ f = tempfile.NamedTemporaryFile(mode="w+b", prefix=prefix or "", suffix=suffix,
1694+ dir=dir, delete=delete)
1695+ else:
1696+ if PY2:
1697+ # Py2 temp files don't accept encoding/newline; writes of unicode will be encoded
1698+ # using the default encoding. If you need strict control, wrap with codecs/io.open.
1699+ f = tempfile.NamedTemporaryFile(mode="w+", prefix=prefix or "", suffix=suffix,
1700+ dir=dir, delete=delete)
1701+ else:
1702+ f = tempfile.NamedTemporaryFile(mode="w+", prefix=prefix or "", suffix=suffix,
1703+ dir=dir, delete=delete, encoding=encoding, newline=newline)
1704+
16121705 if init is not None:
16131706 f.write(init)
16141707 f.seek(0)
0 commit comments