From bea3d02ca7f9dae2b6b31d725fd9665b45fb91b3 Mon Sep 17 00:00:00 2001
From: Jonathan Slenders <jonathan@slenders.be>
Date: Thu, 10 Apr 2025 22:21:44 +0000
Subject: [PATCH] Typecheck examples.

---
 .github/workflows/test.yaml      |  3 ++-
 examples/ssh-and-telnet-embed.py |  6 +++--
 src/ptpython/repl.py             | 43 ++++++++++++++++++++++++++++++--
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 6d2877b3..d53bfcc1 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -19,7 +19,8 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Type Checking
         run: |
-          uvx --with . mypy src/ptpython
+          uvx --with . mypy src/ptpython/
+          uvx --with . mypy examples/
       - name: Code formatting
         if: ${{ matrix.python-version == '3.13' }}
         run: |
diff --git a/examples/ssh-and-telnet-embed.py b/examples/ssh-and-telnet-embed.py
index 62fa76d9..2b293e6f 100755
--- a/examples/ssh-and-telnet-embed.py
+++ b/examples/ssh-and-telnet-embed.py
@@ -6,6 +6,8 @@
 https://gist.github.com/vxgmichel/7685685b3e5ead04ada4a3ba75a48eef
 """
 
+from __future__ import annotations
+
 import asyncio
 import pathlib
 
@@ -15,7 +17,7 @@
     PromptToolkitSSHServer,
     PromptToolkitSSHSession,
 )
-from prompt_toolkit.contrib.telnet.server import TelnetServer
+from prompt_toolkit.contrib.telnet.server import TelnetConnection, TelnetServer
 
 from ptpython.repl import embed
 
@@ -28,7 +30,7 @@ def ensure_key(filename: str = "ssh_host_key") -> str:
     return str(path)
 
 
-async def interact(connection: PromptToolkitSSHSession) -> None:
+async def interact(connection: PromptToolkitSSHSession | TelnetConnection) -> None:
     global_dict = {**globals(), "print": print_formatted_text}
     await embed(return_asyncio_coroutine=True, globals=global_dict)
 
diff --git a/src/ptpython/repl.py b/src/ptpython/repl.py
index ba6717fb..469ed694 100644
--- a/src/ptpython/repl.py
+++ b/src/ptpython/repl.py
@@ -20,7 +20,17 @@
 import warnings
 from dis import COMPILER_FLAG_NAMES
 from pathlib import Path
-from typing import Any, Callable, ContextManager, Iterable, NoReturn, Sequence
+from typing import (
+    Any,
+    Callable,
+    ContextManager,
+    Coroutine,
+    Iterable,
+    Literal,
+    NoReturn,
+    Sequence,
+    overload,
+)
 
 from prompt_toolkit.formatted_text import OneStyleAndTextTuple
 from prompt_toolkit.patch_stdout import patch_stdout as patch_stdout_context
@@ -505,6 +515,34 @@ class ReplExit(Exception):
     """
 
 
+@overload
+def embed(
+    globals: dict[str, Any] | None = ...,
+    locals: dict[str, Any] | None = ...,
+    configure: Callable[[PythonRepl], None] | None = ...,
+    vi_mode: bool = ...,
+    history_filename: str | None = ...,
+    title: str | None = ...,
+    startup_paths: Sequence[str | Path] | None = ...,
+    patch_stdout: bool = ...,
+    return_asyncio_coroutine: Literal[False] = ...,
+) -> None: ...
+
+
+@overload
+def embed(
+    globals: dict[str, Any] | None = ...,
+    locals: dict[str, Any] | None = ...,
+    configure: Callable[[PythonRepl], None] | None = ...,
+    vi_mode: bool = ...,
+    history_filename: str | None = ...,
+    title: str | None = ...,
+    startup_paths: Sequence[str | Path] | None = ...,
+    patch_stdout: bool = ...,
+    return_asyncio_coroutine: Literal[True] = ...,
+) -> Coroutine[Any, Any, None]: ...
+
+
 def embed(
     globals: dict[str, Any] | None = None,
     locals: dict[str, Any] | None = None,
@@ -515,7 +553,7 @@ def embed(
     startup_paths: Sequence[str | Path] | None = None,
     patch_stdout: bool = False,
     return_asyncio_coroutine: bool = False,
-) -> None:
+) -> None | Coroutine[Any, Any, None]:
     """
     Call this to embed  Python shell at the current point in your program.
     It's similar to `IPython.embed` and `bpython.embed`. ::
@@ -577,3 +615,4 @@ async def coroutine() -> None:
     else:
         with patch_context:
             repl.run()
+        return None