From a3ad21144bf8802dc61aed31a176214fac58210f Mon Sep 17 00:00:00 2001
From: Samuel Searles-Bryant <sam@samueljsb.co.uk>
Date: Wed, 15 May 2024 16:52:19 +0100
Subject: [PATCH] Overload hash generation functions to make return types more
 specific

Before this change, the return type for these functions was always
`bytes| str`. However, the type can be narrowed based on the value
passed to `raw_output`.

This change adds overload annotations for the two allowable values to
allow type checkers to narrow the type and correctly infer the exact
return type.

Resolves #2046
---
 faker/providers/misc/__init__.py | 26 +++++++++++++++++-
 faker/proxy.pyi                  | 46 +++++++++++++++++++++++++++++---
 2 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/faker/providers/misc/__init__.py b/faker/providers/misc/__init__.py
index 240249cbba..5fc6f8b131 100644
--- a/faker/providers/misc/__init__.py
+++ b/faker/providers/misc/__init__.py
@@ -9,7 +9,7 @@
 import uuid
 import zipfile
 
-from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple, Type, Union
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Set, Tuple, Type, Union, overload
 
 from faker.exceptions import UnsupportedFeature
 
@@ -56,6 +56,14 @@ def binary(self, length: int = (1 * 1024 * 1024)) -> bytes:
         # Generator is unseeded anyway, just use urandom
         return os.urandom(length)
 
+    @overload
+    def md5(self, raw_output: Literal[True]) -> bytes:
+        ...
+
+    @overload
+    def md5(self, raw_output: Literal[False]) -> str:
+        ...
+
     def md5(self, raw_output: bool = False) -> Union[bytes, str]:
         """Generate a random MD5 hash.
 
@@ -70,6 +78,14 @@ def md5(self, raw_output: bool = False) -> Union[bytes, str]:
             return res.digest()
         return res.hexdigest()
 
+    @overload
+    def sha1(self, raw_output: Literal[True]) -> bytes:
+        ...
+
+    @overload
+    def sha1(self, raw_output: Literal[False]) -> str:
+        ...
+
     def sha1(self, raw_output: bool = False) -> Union[bytes, str]:
         """Generate a random SHA-1 hash.
 
@@ -84,6 +100,14 @@ def sha1(self, raw_output: bool = False) -> Union[bytes, str]:
             return res.digest()
         return res.hexdigest()
 
+    @overload
+    def sha256(self, raw_output: Literal[True]) -> bytes:
+        ...
+
+    @overload
+    def sha256(self, raw_output: Literal[False]) -> str:
+        ...
+
     def sha256(self, raw_output: bool = False) -> Union[bytes, str]:
         """Generate a random SHA-256 hash.
 
diff --git a/faker/proxy.pyi b/faker/proxy.pyi
index 679c9ca11f..1c7e9cf4ed 100644
--- a/faker/proxy.pyi
+++ b/faker/proxy.pyi
@@ -20,6 +20,7 @@ from typing import (
     Type,
     TypeVar,
     Union,
+    overload,
 )
 from uuid import UUID
 
@@ -1944,7 +1945,20 @@ class Faker:
         :meth:`json() <faker.providers.misc.Provider.json>` which is used under the hood.
         """
         ...
-    def md5(self, raw_output: bool = ...) -> Union[bytes, str]:
+    @overload
+    def md5(self, raw_output: Literal[True]) -> bytes:
+        """
+        Generate a random MD5 hash.
+
+        If ``raw_output`` is ``False`` (default), a hexadecimal string representation of the MD5 hash
+        will be returned. If ``True``, a ``bytes`` object representation will be returned instead.
+
+        :sample: raw_output=False
+        :sample: raw_output=True
+        """
+        ...
+    @overload
+    def md5(self, raw_output: Literal[False]) -> str:
         """
         Generate a random MD5 hash.
 
@@ -2001,7 +2015,20 @@ class Faker:
                 num_rows=10, include_row_ids=True
         """
         ...
-    def sha1(self, raw_output: bool = ...) -> Union[bytes, str]:
+    @overload
+    def sha1(self, raw_output: Literal[True]) -> bytes:
+        """
+        Generate a random SHA-1 hash.
+
+        If ``raw_output`` is ``False`` (default), a hexadecimal string representation of the SHA-1 hash
+        will be returned. If ``True``, a ``bytes`` object representation will be returned instead.
+
+        :sample: raw_output=False
+        :sample: raw_output=True
+        """
+        ...
+    @overload
+    def sha1(self, raw_output: Literal[False]) -> str:
         """
         Generate a random SHA-1 hash.
 
@@ -2012,7 +2039,20 @@ class Faker:
         :sample: raw_output=True
         """
         ...
-    def sha256(self, raw_output: bool = ...) -> Union[bytes, str]:
+    @overload
+    def sha256(self, raw_output: Literal[True]) -> bytes:
+        """
+        Generate a random SHA-256 hash.
+
+        If ``raw_output`` is ``False`` (default), a hexadecimal string representation of the SHA-256 hash
+        will be returned. If ``True``, a ``bytes`` object representation will be returned instead.
+
+        :sample: raw_output=False
+        :sample: raw_output=True
+        """
+        ...
+    @overload
+    def sha256(self, raw_output: Literal[False]) -> str:
         """
         Generate a random SHA-256 hash.