diff --git a/colorama/ansi.py b/colorama/ansi.py index 543f39b..26f9e42 100644 --- a/colorama/ansi.py +++ b/colorama/ansi.py @@ -96,7 +96,15 @@ class AnsiStyle(AnsiCodes): NORMAL = 22 RESET_ALL = 0 -Fore = AnsiFore() -Back = AnsiBack() -Style = AnsiStyle() + +def _make_ansi_codes_type(codes): + """Return a type whose attributes are the rendered values from *codes*.""" + attributes = vars(codes).copy() + attributes['__module__'] = __name__ + return type(codes.__class__.__name__, (), attributes) + + +Fore = _make_ansi_codes_type(AnsiFore()) +Back = _make_ansi_codes_type(AnsiBack()) +Style = _make_ansi_codes_type(AnsiStyle()) Cursor = AnsiCursor() diff --git a/colorama/tests/ansi_test.py b/colorama/tests/ansi_test.py index 0a20c80..f15aec7 100644 --- a/colorama/tests/ansi_test.py +++ b/colorama/tests/ansi_test.py @@ -1,5 +1,6 @@ # Copyright Jonathan Hartley 2013. BSD 3-Clause license, see LICENSE file. import sys +from typing import Optional from unittest import TestCase, main from ..ansi import Back, Fore, Style @@ -66,6 +67,15 @@ def testBackAttributes(self): self.assertEqual(Back.LIGHTWHITE_EX, '\033[107m') + def testAnsiCodesCanBeUsedAsTypingTypes(self): + self.assertIsInstance(Back, type) + + def colored_background(color: Optional[Back] = Back.RED): + return color + + self.assertEqual(colored_background(), Back.RED) + + def testStyleAttributes(self): self.assertEqual(Style.DIM, '\033[2m') self.assertEqual(Style.NORMAL, '\033[22m')