From a3d7b94c77411405d058ad517d8414391730b629 Mon Sep 17 00:00:00 2001
From: Ben-Reg <3612364+Ben-Reg@users.noreply.github.com>
Date: Mon, 3 Jan 2022 22:11:08 -0600
Subject: [PATCH 1/4] WIP: Adding type hints to formatter.py

---
 bpython/formatter.py | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/bpython/formatter.py b/bpython/formatter.py
index 9618979a..0458f1e4 100644
--- a/bpython/formatter.py
+++ b/bpython/formatter.py
@@ -24,7 +24,11 @@
 # Pygments really kicks ass, it made it really easy to
 # get the exact behaviour I wanted, thanks Pygments.:)
 
+# mypy: disallow_untyped_defs=True
+# mypy: disallow_untyped_calls=True
 
+
+from typing import MutableMapping, Any, Iterable, TextIO
 from pygments.formatter import Formatter
 from pygments.token import (
     Keyword,
@@ -96,7 +100,11 @@ class BPythonFormatter(Formatter):
     See the Pygments source for more info; it's pretty
     straightforward."""
 
-    def __init__(self, color_scheme, **options):
+    def __init__(
+        self,
+        color_scheme: MutableMapping[str, str],
+        **options: str
+    ) -> None:
         self.f_strings = {}
         for k, v in theme_map.items():
             self.f_strings[k] = f"\x01{color_scheme[v]}"
@@ -106,8 +114,12 @@ def __init__(self, color_scheme, **options):
                 self.f_strings[k] += "I"
         super().__init__(**options)
 
-    def format(self, tokensource, outfile):
-        o = ""
+    def format(
+        self,
+        tokensource: Iterable[MutableMapping[Any, str]],
+        outfile: TextIO
+    ) -> None:
+        o: str = ""
         for token, text in tokensource:
             if text == "\n":
                 continue

From 895fdc0a046693bea17e7dd0e8009828ff780285 Mon Sep 17 00:00:00 2001
From: Ben-Reg <3612364+Ben-Reg@users.noreply.github.com>
Date: Thu, 6 Jan 2022 10:15:11 -0600
Subject: [PATCH 2/4] Formatted with Black

---
 bpython/formatter.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/bpython/formatter.py b/bpython/formatter.py
index 0458f1e4..8aab27e1 100644
--- a/bpython/formatter.py
+++ b/bpython/formatter.py
@@ -101,9 +101,7 @@ class BPythonFormatter(Formatter):
     straightforward."""
 
     def __init__(
-        self,
-        color_scheme: MutableMapping[str, str],
-        **options: str
+        self, color_scheme: MutableMapping[str, str], **options: str
     ) -> None:
         self.f_strings = {}
         for k, v in theme_map.items():
@@ -115,9 +113,7 @@ def __init__(
         super().__init__(**options)
 
     def format(
-        self,
-        tokensource: Iterable[MutableMapping[Any, str]],
-        outfile: TextIO
+        self, tokensource: Iterable[MutableMapping[Any, str]], outfile: TextIO
     ) -> None:
         o: str = ""
         for token, text in tokensource:

From 8d71839cb5e3a54f17de308d1b450cf7d7a1cd81 Mon Sep 17 00:00:00 2001
From: Ben-Reg <3612364+Ben-Reg@users.noreply.github.com>
Date: Thu, 6 Jan 2022 16:14:26 -0600
Subject: [PATCH 3/4] Added temp variable to allow for stricter typing.

---
 bpython/formatter.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

diff --git a/bpython/formatter.py b/bpython/formatter.py
index 8aab27e1..8dade066 100644
--- a/bpython/formatter.py
+++ b/bpython/formatter.py
@@ -28,9 +28,10 @@
 # mypy: disallow_untyped_calls=True
 
 
-from typing import MutableMapping, Any, Iterable, TextIO
+from typing import MutableMapping, Iterable, TextIO, Optional
 from pygments.formatter import Formatter
 from pygments.token import (
+    _TokenType,
     Keyword,
     Name,
     Comment,
@@ -113,16 +114,20 @@ def __init__(
         super().__init__(**options)
 
     def format(
-        self, tokensource: Iterable[MutableMapping[Any, str]], outfile: TextIO
+        self,
+        tokensource: Iterable[MutableMapping[_TokenType, str]],
+        outfile: TextIO,
     ) -> None:
         o: str = ""
         for token, text in tokensource:
+            t: Optional[_TokenType] = token
             if text == "\n":
                 continue
 
-            while token not in self.f_strings:
-                token = token.parent
-            o += f"{self.f_strings[token]}\x03{text}\x04"
+            while t not in self.f_strings:
+                if t is not None:
+                    t = t.parent
+            o += f"{self.f_strings[t]}\x03{text}\x04"
         outfile.write(o.rstrip())
 
 

From 0ca23b1a6de9bc46f74171ec340677a5d261aabb Mon Sep 17 00:00:00 2001
From: Ben-Reg <3612364+Ben-Reg@users.noreply.github.com>
Date: Thu, 6 Jan 2022 17:57:46 -0600
Subject: [PATCH 4/4] Second attempt at stricter typing.

---
 bpython/formatter.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/bpython/formatter.py b/bpython/formatter.py
index 8dade066..20573bd7 100644
--- a/bpython/formatter.py
+++ b/bpython/formatter.py
@@ -28,7 +28,7 @@
 # mypy: disallow_untyped_calls=True
 
 
-from typing import MutableMapping, Iterable, TextIO, Optional
+from typing import MutableMapping, Iterable, TextIO
 from pygments.formatter import Formatter
 from pygments.token import (
     _TokenType,
@@ -120,14 +120,15 @@ def format(
     ) -> None:
         o: str = ""
         for token, text in tokensource:
-            t: Optional[_TokenType] = token
             if text == "\n":
                 continue
 
-            while t not in self.f_strings:
-                if t is not None:
-                    t = t.parent
-            o += f"{self.f_strings[t]}\x03{text}\x04"
+            while token not in self.f_strings:
+                if token.parent is None:
+                    break
+                else:
+                    token = token.parent
+            o += f"{self.f_strings[token]}\x03{text}\x04"
         outfile.write(o.rstrip())