From c5dd0e9f0cba6b9aee637ed9b695179c94f716fc Mon Sep 17 00:00:00 2001
From: Ilan Schnell <ilanschnell@gmail.com>
Date: Thu, 10 Jun 2021 23:40:25 -0500
Subject: [PATCH] add Int1 and Endian types

---
 bitarray/__init__.pyi | 21 +++++++++++----------
 bitarray/util.pyi     | 28 +++++++++++++---------------
 2 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/bitarray/__init__.pyi b/bitarray/__init__.pyi
index 994f7b52..15a08f90 100644
--- a/bitarray/__init__.pyi
+++ b/bitarray/__init__.pyi
@@ -1,11 +1,12 @@
 from collections.abc import Iterable, Iterator
 from unittest.runner import TextTestResult
 
-from typing import Any, BinaryIO, Union, overload
-from typing_extensions import Literal
+from typing import Any, BinaryIO, Literal, Union, overload
 
 
-Bool = Union[Literal[0, 1], bool]
+Int1 = Literal[0, 1]
+Bool = Union[Int1, bool]
+Endian = Literal["little", "big"]
 Codedict = dict[Any, bitarray]
 
 
@@ -18,7 +19,7 @@ class decodetree:
 class bitarray:
     def __init__(self,
                  initializer: Union[int, str, Iterable[Bool], None] = ...,
-                 endian: str = ..., /) -> None: ...
+                 endian: Endian = ..., /) -> None: ...
 
     def all(self) -> bool: ...
     def any(self) -> bool: ...
@@ -34,7 +35,7 @@ class bitarray:
 
     def decode(self, code: Union[Codedict, decodetree], /) -> list: ...
     def encode(self, code: Codedict, x: Iterable, /) -> None: ...
-    def endian(self) -> str: ...
+    def endian(self) -> Endian: ...
     def extend(self, x: Iterable[Bool], /) -> None: ...
     def fill(self) -> int: ...
     def find(self,
@@ -56,7 +57,7 @@ class bitarray:
 
     def itersearch(self, a: Union[bitarray, Bool], /) -> Iterator[int]: ...
     def pack(self, b: bytes, /) -> None: ...
-    def pop(self, i: int = ..., /) -> Bool: ...
+    def pop(self, i: int = ..., /) -> Int1: ...
     def remove(self, value: Bool, /) -> None: ...
     def reverse(self) -> None: ...
     def search(self, a: Union[bitarray, Bool],
@@ -73,9 +74,9 @@ class bitarray:
                one: bytes = ...) -> bytes: ...
 
     def __len__(self) -> int: ...
-    def __iter__(self) -> Iterator[Bool]: ...
+    def __iter__(self) -> Iterator[Int1]: ...
     @overload
-    def __getitem__(self, i: int) -> Bool: ...
+    def __getitem__(self, i: int) -> Int1: ...
     @overload
     def __getitem__(self, s: slice) -> bitarray: ...
     @overload
@@ -113,7 +114,7 @@ class frozenbitarray(bitarray):
 
 
 def bits2bytes(n: int, /) -> int: ...
-def get_default_endian() -> str: ...
+def get_default_endian() -> Endian: ...
 def test(verbosity: int = ..., repeat: int = ...) -> TextTestResult: ...
-def _set_default_endian(endian: str) -> None: ...
+def _set_default_endian(endian: Endian) -> None: ...
 def _sysinfo() -> tuple: ...
diff --git a/bitarray/util.pyi b/bitarray/util.pyi
index 3d7a3905..93a208e3 100644
--- a/bitarray/util.pyi
+++ b/bitarray/util.pyi
@@ -1,24 +1,21 @@
-from typing import Any, AnyStr, BinaryIO, Hashable, Literal, Optional, Union
+from typing import Any, AnyStr, BinaryIO, Hashable, Optional, Union
 
-from bitarray import bitarray
+from bitarray import bitarray, Int1, Bool, Endian
 
 
-Bool = Union[Literal[0, 1], bool]
-
-
-def zeros(length: int, endian: Optional[str] = ...) -> bitarray: ...
-def urandom(length: int, endian: Optional[str] = ...) -> bitarray: ...
+def zeros(length: int, endian: Optional[Endian] = ...) -> bitarray: ...
+def urandom(length: int, endian: Optional[Endian] = ...) -> bitarray: ...
 def pprint(a: Any, stream: BinaryIO = ...,
            group: int = ...,
            indent: int = ...,
            width: int = ...) -> None: ...
 
-def make_endian(a: bitarray, endian: str) -> bitarray: ...
+def make_endian(a: bitarray, endian: Endian) -> bitarray: ...
 def rindex(a: bitarray, Value: Bool) -> int: ...
 def strip(a: bitarray, mode: str = ...) -> bitarray: ...
 
 def count_n(a: bitarray, n: int, /) -> int: ...
-def parity(a: bitarray, /) -> Bool: ...
+def parity(a: bitarray, /) -> Int1: ...
 def count_and(a: bitarray, b: bitarray, /) -> int: ...
 def count_or(a: bitarray, b: bitarray, /) -> int: ...
 def count_xor(a: bitarray, b: bitarray, /) -> int: ...
@@ -26,20 +23,21 @@ def subset(a: bitarray, b: bitarray, /) -> bool: ...
 
 def ba2hex(a: bitarray, /) -> str: ...
 def hex2ba(s: AnyStr, /,
-           endian: Optional[str] = ...) -> bitarray: ...
+           endian: Optional[Endian] = ...) -> bitarray: ...
 
 def ba2base(n: int, a: bitarray, /) -> str: ...
 def base2ba(n: int,
             s: AnyStr, /,
-            endian: Optional[str] = ...) -> bitarray: ...
+            endian: Optional[Endian] = ...) -> bitarray: ...
 
-def ba2int(a: bitarray, signed: bool = ...) -> int: ...
+def ba2int(a: bitarray, signed: int = ...) -> int: ...
 def int2ba(i: int, /,
            length: int = ...,
-           endian: str = ...,
-           signed: bool = ...) -> bitarray: ...
+           endian: Endian = ...,
+           signed: int = ...) -> bitarray: ...
 
 def serialize(a: bitarray, /) -> bytes: ...
 def deserialize(b: bytes, /) -> bitarray: ...
 def huffman_code(freq_map: dict[Hashable, Union[int, float]],
-                 endian: Optional[str] = ...) -> dict[Hashable, bitarray]: ...
+                 endian: Optional[Endian] = ...
+                 ) -> dict[Hashable, bitarray]: ...