Skip to content

Commit

Permalink
Minor type hinting fixes.
Browse files Browse the repository at this point in the history
  • Loading branch information
anselor committed Sep 11, 2020
1 parent 872da20 commit 92b8a38
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 7 deletions.
21 changes: 14 additions & 7 deletions cmd2/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,9 +11,10 @@
import subprocess
import sys
import threading

import unicodedata
from enum import Enum
from typing import Any, Callable, Dict, Iterable, List, Optional, TextIO, Union
from typing import Any, Callable, Dict, IO, Iterable, List, Optional, TextIO, Type, Union

from . import constants

Expand Down Expand Up @@ -470,10 +471,15 @@ def getbytes(self) -> bytes:
"""Get the internal contents as bytes"""
return bytes(self.buffer.byte_buf)

def read(self) -> str:
def read(self, size: Optional[int] = -1) -> str:
"""Read from the internal contents as a str and then clear them out"""
result = self.getvalue()
self.clear()
if size is None or size == -1:
result = self.getvalue()
self.clear()
else:
result = self.buffer.byte_buf[:size].decode(encoding=self.encoding, errors=self.errors)
self.buffer.byte_buf = self.buffer.byte_buf[size:]

return result

def readbytes(self) -> bytes:
Expand Down Expand Up @@ -668,7 +674,7 @@ def __exit__(self, *args) -> None:

class RedirectionSavedState:
"""Created by each command to store information required to restore state after redirection"""
def __init__(self, self_stdout: Union[StdSim, TextIO], sys_stdout: Union[StdSim, TextIO],
def __init__(self, self_stdout: Union[StdSim, IO[str]], sys_stdout: Union[StdSim, IO[str]],
pipe_proc_reader: Optional[ProcReader], saved_redirecting: bool) -> None:
"""
RedirectionSavedState initializer
Expand Down Expand Up @@ -1025,11 +1031,12 @@ def categorize(func: Union[Callable, Iterable[Callable]], category: str) -> None
:Example:
>>> import cmd2
>>> class MyApp(cmd2.Cmd):
>>> def do_echo(self, arglist):
>>> self.poutput(' '.join(arglist)
>>>
>>> utils.categorize(do_echo, "Text Processing")
>>> cmd2.utils.categorize(do_echo, "Text Processing")
For an alternative approach to categorizing commands using a decorator, see
:func:`~cmd2.decorators.with_category`
Expand All @@ -1044,7 +1051,7 @@ def categorize(func: Union[Callable, Iterable[Callable]], category: str) -> None
setattr(func, constants.CMD_ATTR_HELP_CATEGORY, category)


def get_defining_class(meth):
def get_defining_class(meth) -> Type:
"""
Attempts to resolve the class that defined a method.
Expand Down
7 changes: 7 additions & 0 deletions tests/test_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,13 @@ def test_stdsim_read(stdout_sim):
assert stdout_sim.read() == my_str
assert stdout_sim.getvalue() == ''

stdout_sim.write(my_str)

assert stdout_sim.getvalue() == my_str
assert stdout_sim.read(2) == my_str[:2]
assert stdout_sim.getvalue() == my_str[2:]


def test_stdsim_read_bytes(stdout_sim):
b_str = b'Hello World'
stdout_sim.buffer.write(b_str)
Expand Down

0 comments on commit 92b8a38

Please sign in to comment.