diff --git a/plotly/express/__init__.pyi b/plotly/express/__init__.pyi new file mode 100644 index 0000000000..afee65b1ec --- /dev/null +++ b/plotly/express/__init__.pyi @@ -0,0 +1,539 @@ +"""Type stubs for plotly.express - enables IDE autocomplete and static type checking.""" + +from typing import Any, Literal, Sequence, Mapping, Optional, Union +import pandas as pd +import plotly.graph_objs as go + +# Type aliases for common parameter types +DataFrameLike = Union[pd.DataFrame, Mapping[str, Any], None] +ColumnRef = Union[str, int, Sequence[str], Sequence[int], None] +ColorScale = Union[str, Sequence[str], Sequence[tuple[float, str]], None] +SequenceLike = Union[Sequence[Any], None] +MappingLike = Union[Mapping[Any, Any], None] +RangeLike = Union[Sequence[float], None] + +# Chart creation functions + +def scatter( + data_frame: DataFrameLike = None, + x: ColumnRef = None, + y: ColumnRef = None, + color: ColumnRef = None, + symbol: ColumnRef = None, + size: ColumnRef = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + text: ColumnRef = None, + facet_row: ColumnRef = None, + facet_col: ColumnRef = None, + facet_col_wrap: int = 0, + facet_row_spacing: Optional[float] = None, + facet_col_spacing: Optional[float] = None, + error_x: ColumnRef = None, + error_x_minus: ColumnRef = None, + error_y: ColumnRef = None, + error_y_minus: ColumnRef = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + orientation: Optional[Literal["v", "h"]] = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + color_continuous_scale: ColorScale = None, + range_color: RangeLike = None, + color_continuous_midpoint: Optional[float] = None, + symbol_sequence: SequenceLike = None, + symbol_map: MappingLike = None, + opacity: Optional[float] = None, + size_max: Optional[float] = None, + marginal_x: Optional[Literal["rug", "box", "violin", "histogram"]] = None, + marginal_y: Optional[Literal["rug", "box", "violin", "histogram"]] = None, + trendline: Optional[Literal["ols", "lowess", "rolling", "expanding", "ewm"]] = None, + trendline_options: MappingLike = None, + trendline_color_override: Optional[str] = None, + trendline_scope: Literal["trace", "overall"] = "trace", + log_x: bool = False, + log_y: bool = False, + range_x: RangeLike = None, + range_y: RangeLike = None, + render_mode: Literal["auto", "svg", "webgl"] = "auto", + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... +def line( + data_frame: DataFrameLike = None, + x: ColumnRef = None, + y: ColumnRef = None, + line_group: ColumnRef = None, + color: ColumnRef = None, + line_dash: ColumnRef = None, + symbol: ColumnRef = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + text: ColumnRef = None, + facet_row: ColumnRef = None, + facet_col: ColumnRef = None, + facet_col_wrap: int = 0, + facet_row_spacing: Optional[float] = None, + facet_col_spacing: Optional[float] = None, + error_x: ColumnRef = None, + error_x_minus: ColumnRef = None, + error_y: ColumnRef = None, + error_y_minus: ColumnRef = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + orientation: Optional[Literal["v", "h"]] = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + line_dash_sequence: SequenceLike = None, + line_dash_map: MappingLike = None, + symbol_sequence: SequenceLike = None, + symbol_map: MappingLike = None, + markers: bool = False, + log_x: bool = False, + log_y: bool = False, + range_x: RangeLike = None, + range_y: RangeLike = None, + line_shape: Optional[Literal["linear", "spline", "hv", "vh", "hvh", "vhv"]] = None, + render_mode: Literal["auto", "svg", "webgl"] = "auto", + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... +def bar( + data_frame: DataFrameLike = None, + x: ColumnRef = None, + y: ColumnRef = None, + color: ColumnRef = None, + pattern_shape: ColumnRef = None, + facet_row: ColumnRef = None, + facet_col: ColumnRef = None, + facet_col_wrap: int = 0, + facet_row_spacing: Optional[float] = None, + facet_col_spacing: Optional[float] = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + text: ColumnRef = None, + base: ColumnRef = None, + error_x: ColumnRef = None, + error_x_minus: ColumnRef = None, + error_y: ColumnRef = None, + error_y_minus: ColumnRef = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + color_continuous_scale: ColorScale = None, + pattern_shape_sequence: SequenceLike = None, + pattern_shape_map: MappingLike = None, + range_color: RangeLike = None, + color_continuous_midpoint: Optional[float] = None, + opacity: Optional[float] = None, + orientation: Optional[Literal["v", "h"]] = None, + barmode: Literal["relative", "overlay", "group"] = "relative", + log_x: bool = False, + log_y: bool = False, + range_x: RangeLike = None, + range_y: RangeLike = None, + text_auto: Union[bool, str] = False, + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... +def histogram( + data_frame: DataFrameLike = None, + x: ColumnRef = None, + y: ColumnRef = None, + color: ColumnRef = None, + pattern_shape: ColumnRef = None, + facet_row: ColumnRef = None, + facet_col: ColumnRef = None, + facet_col_wrap: int = 0, + facet_row_spacing: Optional[float] = None, + facet_col_spacing: Optional[float] = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + pattern_shape_sequence: SequenceLike = None, + pattern_shape_map: MappingLike = None, + marginal: Optional[Literal["rug", "box", "violin", "histogram"]] = None, + opacity: Optional[float] = None, + orientation: Optional[Literal["v", "h"]] = None, + barmode: Literal["relative", "overlay", "group"] = "relative", + barnorm: Optional[Literal["", "fraction", "percent"]] = None, + histnorm: Optional[ + Literal["", "percent", "probability", "density", "probability density"] + ] = None, + log_x: bool = False, + log_y: bool = False, + range_x: RangeLike = None, + range_y: RangeLike = None, + histfunc: Optional[Literal["count", "sum", "avg", "min", "max"]] = None, + cumulative: Optional[bool] = None, + nbins: Optional[int] = None, + text_auto: Union[bool, str] = False, + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... +def box( + data_frame: DataFrameLike = None, + x: ColumnRef = None, + y: ColumnRef = None, + color: ColumnRef = None, + facet_row: ColumnRef = None, + facet_col: ColumnRef = None, + facet_col_wrap: int = 0, + facet_row_spacing: Optional[float] = None, + facet_col_spacing: Optional[float] = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + orientation: Optional[Literal["v", "h"]] = None, + boxmode: Optional[Literal["group", "overlay"]] = None, + log_x: bool = False, + log_y: bool = False, + range_x: RangeLike = None, + range_y: RangeLike = None, + points: Optional[Literal["outliers", "suspectedoutliers", "all", False]] = None, + notched: bool = False, + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... +def violin( + data_frame: DataFrameLike = None, + x: ColumnRef = None, + y: ColumnRef = None, + color: ColumnRef = None, + facet_row: ColumnRef = None, + facet_col: ColumnRef = None, + facet_col_wrap: int = 0, + facet_row_spacing: Optional[float] = None, + facet_col_spacing: Optional[float] = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + orientation: Optional[Literal["v", "h"]] = None, + violinmode: Optional[Literal["group", "overlay"]] = None, + log_x: bool = False, + log_y: bool = False, + range_x: RangeLike = None, + range_y: RangeLike = None, + points: Optional[Literal["outliers", "suspectedoutliers", "all", False]] = None, + box: bool = False, + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... +def area( + data_frame: DataFrameLike = None, + x: ColumnRef = None, + y: ColumnRef = None, + line_group: ColumnRef = None, + color: ColumnRef = None, + pattern_shape: ColumnRef = None, + symbol: ColumnRef = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + text: ColumnRef = None, + facet_row: ColumnRef = None, + facet_col: ColumnRef = None, + facet_col_wrap: int = 0, + facet_row_spacing: Optional[float] = None, + facet_col_spacing: Optional[float] = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + pattern_shape_sequence: SequenceLike = None, + pattern_shape_map: MappingLike = None, + symbol_sequence: SequenceLike = None, + symbol_map: MappingLike = None, + markers: bool = False, + orientation: Optional[Literal["v", "h"]] = None, + groupnorm: Optional[Literal["", "fraction", "percent"]] = None, + log_x: bool = False, + log_y: bool = False, + range_x: RangeLike = None, + range_y: RangeLike = None, + line_shape: Optional[Literal["linear", "spline", "hv", "vh", "hvh", "vhv"]] = None, + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... +def pie( + data_frame: DataFrameLike = None, + names: ColumnRef = None, + values: ColumnRef = None, + color: ColumnRef = None, + facet_row: ColumnRef = None, + facet_col: ColumnRef = None, + facet_col_wrap: int = 0, + facet_row_spacing: Optional[float] = None, + facet_col_spacing: Optional[float] = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, + opacity: Optional[float] = None, + hole: Optional[float] = None, +) -> go.Figure: ... + +# 3D chart functions + +def scatter_3d( + data_frame: DataFrameLike = None, + x: ColumnRef = None, + y: ColumnRef = None, + z: ColumnRef = None, + color: ColumnRef = None, + symbol: ColumnRef = None, + size: ColumnRef = None, + text: ColumnRef = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + error_x: ColumnRef = None, + error_x_minus: ColumnRef = None, + error_y: ColumnRef = None, + error_y_minus: ColumnRef = None, + error_z: ColumnRef = None, + error_z_minus: ColumnRef = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + size_max: Optional[float] = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + color_continuous_scale: ColorScale = None, + range_color: RangeLike = None, + color_continuous_midpoint: Optional[float] = None, + symbol_sequence: SequenceLike = None, + symbol_map: MappingLike = None, + opacity: Optional[float] = None, + log_x: bool = False, + log_y: bool = False, + log_z: bool = False, + range_x: RangeLike = None, + range_y: RangeLike = None, + range_z: RangeLike = None, + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... +def line_3d( + data_frame: DataFrameLike = None, + x: ColumnRef = None, + y: ColumnRef = None, + z: ColumnRef = None, + color: ColumnRef = None, + line_dash: ColumnRef = None, + text: ColumnRef = None, + line_group: ColumnRef = None, + symbol: ColumnRef = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + error_x: ColumnRef = None, + error_x_minus: ColumnRef = None, + error_y: ColumnRef = None, + error_y_minus: ColumnRef = None, + error_z: ColumnRef = None, + error_z_minus: ColumnRef = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + line_dash_sequence: SequenceLike = None, + line_dash_map: MappingLike = None, + symbol_sequence: SequenceLike = None, + symbol_map: MappingLike = None, + markers: bool = False, + log_x: bool = False, + log_y: bool = False, + log_z: bool = False, + range_x: RangeLike = None, + range_y: RangeLike = None, + range_z: RangeLike = None, + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... + +# Geographic chart functions + +def scatter_geo( + data_frame: DataFrameLike = None, + lat: ColumnRef = None, + lon: ColumnRef = None, + locations: ColumnRef = None, + locationmode: Optional[ + Literal["ISO-3", "USA-states", "country names", "geojson-id"] + ] = None, + color: ColumnRef = None, + text: ColumnRef = None, + facet_row: ColumnRef = None, + facet_col: ColumnRef = None, + facet_col_wrap: int = 0, + facet_row_spacing: Optional[float] = None, + facet_col_spacing: Optional[float] = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + size: ColumnRef = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + color_continuous_scale: ColorScale = None, + range_color: RangeLike = None, + color_continuous_midpoint: Optional[float] = None, + symbol: ColumnRef = None, + symbol_sequence: SequenceLike = None, + symbol_map: MappingLike = None, + size_max: Optional[float] = None, + opacity: Optional[float] = None, + projection: Optional[str] = None, + scope: Optional[ + Literal[ + "world", "usa", "europe", "asia", "africa", "north america", "south america" + ] + ] = None, + center: MappingLike = None, + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... +def line_geo( + data_frame: DataFrameLike = None, + lat: ColumnRef = None, + lon: ColumnRef = None, + locations: ColumnRef = None, + locationmode: Optional[ + Literal["ISO-3", "USA-states", "country names", "geojson-id"] + ] = None, + color: ColumnRef = None, + line_dash: ColumnRef = None, + text: ColumnRef = None, + facet_row: ColumnRef = None, + facet_col: ColumnRef = None, + facet_col_wrap: int = 0, + facet_row_spacing: Optional[float] = None, + facet_col_spacing: Optional[float] = None, + hover_name: ColumnRef = None, + hover_data: Union[SequenceLike, MappingLike] = None, + custom_data: SequenceLike = None, + line_group: ColumnRef = None, + symbol: ColumnRef = None, + animation_frame: ColumnRef = None, + animation_group: ColumnRef = None, + category_orders: MappingLike = None, + labels: MappingLike = None, + color_discrete_sequence: SequenceLike = None, + color_discrete_map: MappingLike = None, + line_dash_sequence: SequenceLike = None, + line_dash_map: MappingLike = None, + symbol_sequence: SequenceLike = None, + symbol_map: MappingLike = None, + markers: bool = False, + projection: Optional[str] = None, + scope: Optional[ + Literal[ + "world", "usa", "europe", "asia", "africa", "north america", "south america" + ] + ] = None, + center: MappingLike = None, + title: Optional[str] = None, + subtitle: Optional[str] = None, + template: Union[str, go.layout.Template, None] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> go.Figure: ... + +# Utility functions + +def set_mapbox_access_token(token: str) -> None: ... +def get_trendline_results(fig: go.Figure) -> Optional[pd.DataFrame]: ... + +# Special input classes +class IdentityMap: ... + +class Constant: + def __init__(self, value: Any) -> None: ... + +class Range: + def __init__(self, start: float, stop: float, step: float = 1.0) -> None: ... + +# Constants +NO_COLOR: str + +# Submodules +data: Any +colors: Any +trendline_functions: Any + +# Note: Additional functions like scatter_polar, scatter_ternary, scatter_mapbox, etc. +# follow similar patterns and can be added as needed. diff --git a/test_type_stubs.py b/test_type_stubs.py new file mode 100644 index 0000000000..61f8e61f2e --- /dev/null +++ b/test_type_stubs.py @@ -0,0 +1,36 @@ +"""Test script to validate plotly.express type stubs.""" + +import plotly.express as px +import plotly.graph_objs as go +import pandas as pd + +# Create test data +df = pd.DataFrame( + { + "x": [1, 2, 3, 4, 5], + "y": [2, 4, 3, 5, 6], + "color": ["A", "B", "A", "B", "A"], + "size": [10, 20, 15, 25, 30], + } +) + +# Test scatter - should type-check correctly +fig1: go.Figure = px.scatter(df, x="x", y="y", color="color", size="size") + +# Test line +fig2: go.Figure = px.line(df, x="x", y="y", color="color") + +# Test bar +fig3: go.Figure = px.bar(df, x="color", y="y") + +# Test histogram +fig4: go.Figure = px.histogram(df, x="y") + +# Test box +fig5: go.Figure = px.box(df, x="color", y="y") + +# Test with wrong types (should fail type check) +# fig_bad: int = px.scatter(df, x='x', y='y') # Uncomment to test type error + +print("✓ All type annotations work correctly!") +print(f"Figure type: {type(fig1)}")