Improve types for pd.options.display.* properties

Author: mike-flowers-airbnb

pandas-stubs/_config/config.pyi

@@ -49,30 +49,30 @@ class DisplayUnicode(DictWrapper):
 
 class Display(DictWrapper):
     chop_threshold: int | None
-    colheader_justify: str
+    colheader_justify: Literal["left", "right"]
     date_dayfirst: bool
     date_yearfirst: bool
     encoding: str
     expand_frame_repr: bool
     float_format: Callable[[float], str] | None
     html: DisplayHTML
-    large_repr: str
+    large_repr: Literal["truncate", "info"]
     latex: DisplayLaTeX
     max_categories: int
-    max_columns: int
-    max_colwidth: int
-    max_dir_items: int
+    max_columns: int | None
+    max_colwidth: int | None
+    max_dir_items: int | None
     max_info_columns: int
     max_info_rows: int
-    max_rows: int
-    max_seq_items: int
-    memory_usage: bool
-    min_rows: int
+    max_rows: int | None
+    max_seq_items: int | None
+    memory_usage: bool | Literal["deep"] | None
+    min_rows: int | None
     multi_sparse: bool
     notebook_repr_html: bool
     pprint_nest_depth: int
     precision: int
-    show_dimensions: str
+    show_dimensions: bool | Literal["truncate"]
     unicode: DisplayUnicode
     width: int
 

tests/test_config.py

@@ -2,6 +2,7 @@
     TYPE_CHECKING,
     Any,
     Callable,
+    Literal,
     Optional,
 )
 
@@ -48,3 +49,78 @@ def test_display_float_format():
     formatter = "{,.2f}".format
     with pd.option_context("display.float_format", formatter):
         assert pd.get_option("display.float_format") == formatter
+
+
+def test_display_types_none_allowed():
+    # GH 1230
+    # Initial values
+    assert_type(pd.options.display.max_columns, Optional[int])
+    assert_type(pd.options.display.max_colwidth, Optional[int])
+    assert_type(pd.options.display.max_dir_items, Optional[int])
+    assert_type(pd.options.display.max_rows, Optional[int])
+    assert_type(pd.options.display.max_seq_items, Optional[int])
+    assert_type(pd.options.display.min_rows, Optional[int])
+    # Test with None
+    pd.options.display.max_columns = None
+    check(assert_type(pd.options.display.max_columns, None), type(None))
+    pd.options.display.max_colwidth = None
+    check(assert_type(pd.options.display.max_colwidth, None), type(None))
+    pd.options.display.max_dir_items = None
+    check(assert_type(pd.options.display.max_dir_items, None), type(None))
+    pd.options.display.max_rows = None
+    check(assert_type(pd.options.display.max_rows, None), type(None))
+    pd.options.display.max_seq_items = None
+    check(assert_type(pd.options.display.max_seq_items, None), type(None))
+    pd.options.display.min_rows = None
+    check(assert_type(pd.options.display.min_rows, None), type(None))
+    # Test with integer values
+    pd.options.display.max_columns = 100
+    check(assert_type(pd.options.display.max_columns, int), int)
+    pd.options.display.max_colwidth = 100
+    check(assert_type(pd.options.display.max_colwidth, int), int)
+    pd.options.display.max_dir_items = 100
+    check(assert_type(pd.options.display.max_dir_items, int), int)
+    pd.options.display.max_rows = 100
+    check(assert_type(pd.options.display.max_rows, int), int)
+    pd.options.display.max_seq_items = 100
+    check(assert_type(pd.options.display.max_seq_items, int), int)
+    pd.options.display.min_rows = 100
+    check(assert_type(pd.options.display.min_rows, int), int)
+
+
+def test_display_types_literal_constraints():
+    # GH 1230
+    # Various display options have specific allowed values
+    # Test colheader_justify with allowed values
+    assert_type(pd.options.display.colheader_justify, Literal["left", "right"])
+    pd.options.display.colheader_justify = "left"
+    check(assert_type(pd.options.display.colheader_justify, Literal["left"]), str)
+    pd.options.display.colheader_justify = "right"
+    check(assert_type(pd.options.display.colheader_justify, Literal["right"]), str)
+
+    # Test large_repr with allowed values
+    assert_type(pd.options.display.large_repr, Literal["truncate", "info"])
+    pd.options.display.large_repr = "truncate"
+    check(assert_type(pd.options.display.large_repr, Literal["truncate"]), str)
+    pd.options.display.large_repr = "info"
+    check(assert_type(pd.options.display.large_repr, Literal["info"]), str)
+
+    # Test memory_usage with allowed values
+    assert_type(pd.options.display.memory_usage, Optional[Literal[True, False, "deep"]])
+    pd.options.display.memory_usage = True
+    check(assert_type(pd.options.display.memory_usage, Literal[True]), bool)
+    pd.options.display.memory_usage = False
+    check(assert_type(pd.options.display.memory_usage, Literal[False]), bool)
+    pd.options.display.memory_usage = "deep"
+    check(assert_type(pd.options.display.memory_usage, Literal["deep"]), str)
+    pd.options.display.memory_usage = None
+    check(assert_type(pd.options.display.memory_usage, None), type(None))
+
+    # Test show_dimensions with allowed values
+    assert_type(pd.options.display.show_dimensions, Literal[True, False, "truncate"])
+    pd.options.display.show_dimensions = True
+    check(assert_type(pd.options.display.show_dimensions, Literal[True]), bool)
+    pd.options.display.show_dimensions = False
+    check(assert_type(pd.options.display.show_dimensions, Literal[False]), bool)
+    pd.options.display.show_dimensions = "truncate"
+    check(assert_type(pd.options.display.show_dimensions, Literal["truncate"]), str)