Coverage for src/jointview/columns.py: 100%

1"""Which columns the app can draw, and cutting a chosen pair down to its common sample. 

2 

3This is the vocabulary both halves of the app are built on, and it belongs to neither 

4of them. :mod:`jointview.plot` draws what :func:`aligned` produces and 

5:mod:`jointview.stats` summarises it, so the column names and the shaping rules live 

6here rather than in either — two peers reaching into a third, instead of one reaching 

7into the other. 

8 

9``PERIOD`` is the clearest case. It is the name :func:`aligned` writes the x-axis under 

10and the name the statistics read it back out of: a data contract between the two, not 

11a fact about plotting. 

12""" 

13 

14from __future__ import annotations 

15 

16import polars as pl 

17 

18PERIOD = "period" 

19 

20 

21def series_columns(frame: pl.DataFrame) -> list[str]: 

22 """The columns that can be drawn: every numeric one. 

23 

24 >>> import datetime as dt, polars as pl 

25 >>> frame = pl.DataFrame({"date": [dt.date(2024, 1, 1)], "nav": [1.0], "label": ["a"]}) 

26 >>> series_columns(frame) 

27 ['nav'] 

28 """ 

29 return [name for name, dtype in frame.schema.items() if dtype.is_numeric()] 

30 

31 

32def date_column(frame: pl.DataFrame) -> str | None: 

33 """The first temporal column, which becomes the x-axis. None means row number. 

34 

35 >>> import datetime as dt, polars as pl 

36 >>> date_column(pl.DataFrame({"when": [dt.date(2024, 1, 1)], "nav": [1.0]})) 

37 'when' 

38 >>> date_column(pl.DataFrame({"nav": [1.0]})) is None 

39 True 

40 """ 

41 return next((name for name, dtype in frame.schema.items() if dtype.is_temporal()), None) 

42 

43 

44def default_pair(frame: pl.DataFrame) -> tuple[int, int]: 

45 """Indices into :func:`series_columns` to open on — the first two series. 

46 

47 A frame with a single series opens on it twice, rather than refusing to draw. 

48 

49 >>> import polars as pl 

50 >>> default_pair(pl.DataFrame({"a": [1.0], "b": [2.0]})) 

51 (0, 1) 

52 >>> default_pair(pl.DataFrame({"only": [1.0]})) 

53 (0, 0) 

54 """ 

55 names = series_columns(frame) 

56 if not names: 

57 raise ValueError("frame has no numeric columns to plot") # noqa: TRY003 

58 return 0, 1 if len(names) > 1 else 0 

59 

60 

61def aligned(frame: pl.DataFrame, a: str, b: str) -> pl.DataFrame: 

62 """The two series on their common sample: ``period``, ``a``, ``b``, in order. 

63 

64 Both the picture and the summary tables are built from this, so the numbers 

65 beside the chart always describe the lines in it. Renaming also sidesteps 

66 Vega-Lite's field-shorthand escaping and lets ``a`` and ``b`` be the same column. 

67 

68 The three names are fixed whatever the columns were called, the rows come out 

69 sorted by period, and a date where either series is missing is not part of the 

70 sample — here the frame arrives unsorted and with a gap on the 2nd: 

71 

72 >>> import datetime as dt, polars as pl 

73 >>> frame = pl.DataFrame( 

74 ... { 

75 ... "date": [dt.date(2024, 1, 3), dt.date(2024, 1, 1), dt.date(2024, 1, 2)], 

76 ... "x": [3.0, 1.0, None], 

77 ... "y": [30.0, 10.0, 20.0], 

78 ... } 

79 ... ) 

80 >>> pair = aligned(frame, "x", "y") 

81 >>> pair.columns 

82 ['period', 'a', 'b'] 

83 >>> pair["a"].to_list() 

84 [1.0, 3.0] 

85 """ 

86 for column in (a, b): 

87 if column not in frame.columns: 

88 raise KeyError(f"no column {column!r} in frame") # noqa: TRY003 

89 if not frame.schema[column].is_numeric(): 

90 raise TypeError(f"column {column!r} is {frame.schema[column]}, which cannot be drawn") # noqa: TRY003 

91 

92 date = date_column(frame) 

93 period = pl.col(date).alias(PERIOD) if date else pl.int_range(pl.len()).alias(PERIOD) 

94 data = frame.select(period, pl.col(a).alias("a"), pl.col(b).alias("b")) 

95 return data.drop_nulls().sort(PERIOD)