Skip to content

Commit ba787f6

Browse files
subotacsubotac
authored
[hashlib] Fix SHAKE return types for hashlib.new (#16481)
1 parent 1ee8b92 commit ba787f6

3 files changed

Lines changed: 120 additions & 9 deletions

File tree

‎stdlib/@tests/test_cases/check_hashlib.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
from __future__ import annotations
2+
3+
import _hashlib
4+
import hashlib
5+
import sys
6+
from _hashlib import HASH, HASHXOF
7+
from io import BytesIO
8+
from typing_extensions import assert_type
9+
10+
11+
def check_new(name: str) -> None:
12+
for shake in (
13+
hashlib.new("shake_128"),
14+
hashlib.new("shake_256"),
15+
hashlib.new("SHAKE-128", b"data", usedforsecurity=False),
16+
hashlib.new("SHAKE-256"),
17+
):
18+
assert_type(shake, HASHXOF)
19+
assert_type(shake.copy(), HASHXOF)
20+
assert_type(shake.digest(16), bytes)
21+
assert_type(shake.hexdigest(length=16), str)
22+
shake.digest() # type: ignore
23+
shake.hexdigest() # type: ignore
24+
25+
fixed = hashlib.new("sha256")
26+
assert_type(fixed, HASH)
27+
assert_type(fixed.digest(), bytes)
28+
fixed.digest(16) # type: ignore
29+
30+
dynamic = hashlib.new(name)
31+
assert_type(dynamic.digest(), bytes)
32+
assert_type(dynamic.digest(16), bytes)
33+
assert_type(dynamic.copy().hexdigest(length=16), str)
34+
dynamic.digest("16") # type: ignore
35+
36+
assert_type(_hashlib.new("shake_128", string=b"data").digest(16), bytes)
37+
assert_type(_hashlib.new("SHAKE-256").hexdigest(16), str)
38+
assert_type(_hashlib.new("MD5"), HASH)
39+
assert_type(_hashlib.new(name).hexdigest(16), str)
40+
if sys.version_info >= (3, 13):
41+
assert_type(_hashlib.new("shake_256", data=b"data").copy(), HASHXOF)
42+
43+
if sys.version_info >= (3, 11):
44+
hashlib.file_digest(BytesIO(b"data"), lambda: hashlib.new("sha256"))

‎stdlib/_hashlib.pyi‎

Lines changed: 65 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,39 @@ import sys
22
from _typeshed import ReadableBuffer
33
from collections.abc import Callable
44
from types import ModuleType
5-
from typing import AnyStr, Protocol, TypeAlias, final, overload, type_check_only
5+
from typing import AnyStr, Literal, Protocol, TypeAlias, final, overload, type_check_only
66
from typing_extensions import Self, disjoint_base
77

88
_DigestMod: TypeAlias = str | Callable[[], _HashObject] | ModuleType | None
9+
_FixedDigestName: TypeAlias = Literal[
10+
"md5",
11+
"MD5",
12+
"sha1",
13+
"SHA1",
14+
"sha224",
15+
"SHA224",
16+
"sha256",
17+
"SHA256",
18+
"sha384",
19+
"SHA384",
20+
"sha512",
21+
"SHA512",
22+
"sha3_224",
23+
"sha3-224",
24+
"SHA3-224",
25+
"sha3_256",
26+
"sha3-256",
27+
"SHA3-256",
28+
"sha3_384",
29+
"sha3-384",
30+
"SHA3-384",
31+
"sha3_512",
32+
"sha3-512",
33+
"SHA3-512",
34+
]
35+
_XofDigestName: TypeAlias = Literal[
36+
"shake_128", "shake128", "SHAKE128", "shake-128", "SHAKE-128", "shake_256", "shake256", "SHAKE256", "shake-256", "SHAKE-256"
37+
]
938

1039
openssl_md_meth_names: frozenset[str]
1140

@@ -22,17 +51,31 @@ class _HashObject(Protocol):
2251
def hexdigest(self) -> str: ...
2352
def update(self, obj: ReadableBuffer, /) -> None: ...
2453

54+
@type_check_only
55+
class _HashObjectWithOptionalLength(Protocol):
56+
@property
57+
def digest_size(self) -> int: ...
58+
@property
59+
def block_size(self) -> int: ...
60+
@property
61+
def name(self) -> str: ...
62+
def copy(self) -> Self: ...
63+
# Depending on the algorithm, length is either required or not accepted.
64+
def digest(self, length: int = ...) -> bytes: ...
65+
def hexdigest(self, length: int = ...) -> str: ...
66+
def update(self, obj: ReadableBuffer, /) -> None: ...
67+
2568
@disjoint_base
26-
class HASH:
69+
class HASH(_HashObjectWithOptionalLength):
2770
@property
2871
def digest_size(self) -> int: ...
2972
@property
3073
def block_size(self) -> int: ...
3174
@property
3275
def name(self) -> str: ...
3376
def copy(self) -> Self: ...
34-
def digest(self) -> bytes: ...
35-
def hexdigest(self) -> str: ...
77+
def digest(self) -> bytes: ... # type: ignore[override]
78+
def hexdigest(self) -> str: ... # type: ignore[override]
3679
def update(self, obj: ReadableBuffer, /) -> None: ...
3780

3881
class UnsupportedDigestmodError(ValueError): ...
@@ -63,9 +106,19 @@ def get_fips_mode() -> int: ...
63106
def hmac_new(key: ReadableBuffer, msg: ReadableBuffer = b"", digestmod: _DigestMod = None) -> HMAC: ...
64107

65108
if sys.version_info >= (3, 13):
109+
@overload
66110
def new(
67-
name: str, data: ReadableBuffer = b"", *, usedforsecurity: bool = True, string: ReadableBuffer | None = None
111+
name: _XofDigestName, data: ReadableBuffer = b"", *, usedforsecurity: bool = True, string: ReadableBuffer | None = None
112+
) -> HASHXOF: ...
113+
@overload
114+
def new(
115+
name: _FixedDigestName, data: ReadableBuffer = b"", *, usedforsecurity: bool = True, string: ReadableBuffer | None = None
68116
) -> HASH: ...
117+
@overload
118+
def new(
119+
name: str, data: ReadableBuffer = b"", *, usedforsecurity: bool = True, string: ReadableBuffer | None = None
120+
) -> _HashObjectWithOptionalLength: ...
121+
69122
def openssl_md5(
70123
data: ReadableBuffer = b"", *, usedforsecurity: bool = True, string: ReadableBuffer | None = None
71124
) -> HASH: ...
@@ -102,9 +155,14 @@ if sys.version_info >= (3, 13):
102155
def openssl_shake_256(
103156
data: ReadableBuffer = b"", *, usedforsecurity: bool = True, string: ReadableBuffer | None = None
104157
) -> HASHXOF: ...
105-
106158
else:
107-
def new(name: str, string: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> HASH: ...
159+
@overload
160+
def new(name: _XofDigestName, string: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> HASHXOF: ...
161+
@overload
162+
def new(name: _FixedDigestName, string: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> HASH: ...
163+
@overload
164+
def new(name: str, string: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> _HashObjectWithOptionalLength: ...
165+
108166
def openssl_md5(string: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> HASH: ...
109167
def openssl_sha1(string: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> HASH: ...
110168
def openssl_sha224(string: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> HASH: ...

‎stdlib/hashlib.pyi‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,11 @@ import sys
22
from _blake2 import blake2b as blake2b, blake2s as blake2s
33
from _hashlib import (
44
HASH,
5+
HASHXOF,
6+
_FixedDigestName,
57
_HashObject,
8+
_HashObjectWithOptionalLength,
9+
_XofDigestName,
610
openssl_md5 as md5,
711
openssl_sha1 as sha1,
812
openssl_sha3_224 as sha3_224,
@@ -20,7 +24,7 @@ from _hashlib import (
2024
)
2125
from _typeshed import ReadableBuffer
2226
from collections.abc import Callable, Set as AbstractSet
23-
from typing import Protocol, type_check_only
27+
from typing import Protocol, overload, type_check_only
2428

2529
if sys.version_info >= (3, 15):
2630
__all__ = (
@@ -89,7 +93,12 @@ else:
8993
"pbkdf2_hmac",
9094
)
9195

92-
def new(name: str, data: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> HASH: ...
96+
@overload
97+
def new(name: _XofDigestName, data: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> HASHXOF: ...
98+
@overload
99+
def new(name: _FixedDigestName, data: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> HASH: ...
100+
@overload
101+
def new(name: str, data: ReadableBuffer = b"", *, usedforsecurity: bool = True) -> _HashObjectWithOptionalLength: ...
93102

94103
algorithms_guaranteed: AbstractSet[str]
95104
algorithms_available: AbstractSet[str]

0 commit comments

Comments
 (0)