Skip to content

Commit a737f64

Browse files
authored
Merge pull request #493 from palewire/main
refactor(charts): extract annotations to shared mixin
2 parents e474476 + bab4474 commit a737f64

11 files changed

Lines changed: 363 additions & 1454 deletions

File tree

.clinerules

Lines changed: 174 additions & 1210 deletions
Large diffs are not rendered by default.

datawrapper/charts/area.py

Lines changed: 13 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -13,22 +13,25 @@
1313
PlotHeightMode,
1414
)
1515
from .models import (
16+
AnnotationsMixin,
1617
CustomRangeMixin,
1718
CustomTicksMixin,
1819
GridDisplayMixin,
1920
GridFormatMixin,
20-
RangeAnnotation,
21-
TextAnnotation,
2221
)
2322
from .serializers import (
2423
ColorCategory,
25-
ModelListSerializer,
2624
PlotHeight,
2725
)
2826

2927

3028
class AreaChart(
31-
GridDisplayMixin, GridFormatMixin, CustomRangeMixin, CustomTicksMixin, BaseChart
29+
GridDisplayMixin,
30+
GridFormatMixin,
31+
CustomRangeMixin,
32+
CustomTicksMixin,
33+
AnnotationsMixin,
34+
BaseChart,
3235
):
3336
"""A base class for the Datawrapper API's area chart."""
3437

@@ -206,24 +209,6 @@ class AreaChart(
206209
description="The ratio of the plot height",
207210
)
208211

209-
#
210-
# Annotations
211-
#
212-
213-
#: A list of text annotations to display on the chart
214-
text_annotations: list[TextAnnotation | dict[Any, Any]] = Field(
215-
default_factory=list,
216-
alias="text-annotations",
217-
description="A list of text annotations to display on the chart",
218-
)
219-
220-
#: A list of range annotations to display on the chart
221-
range_annotations: list[RangeAnnotation | dict[Any, Any]] = Field(
222-
default_factory=list,
223-
alias="range-annotations",
224-
description="A list of range annotations to display on the chart",
225-
)
226-
227212
@field_validator("interpolation")
228213
@classmethod
229214
def validate_interpolation(
@@ -287,17 +272,14 @@ def serialize_model(self) -> dict:
287272
self.plot_height_fixed,
288273
self.plot_height_ratio,
289274
),
290-
# Annotations
291-
"text-annotations": ModelListSerializer.serialize(
292-
self.text_annotations, TextAnnotation
293-
),
294-
"range-annotations": ModelListSerializer.serialize(
295-
self.range_annotations, RangeAnnotation
296-
),
297275
}
298276

277+
# Add the visualize data to the model
299278
model["metadata"]["visualize"].update(visualize_data)
300279

280+
# Add annotations from mixin
281+
model["metadata"]["visualize"].update(self._serialize_annotations())
282+
301283
# Return the serialized data
302284
return model
303285

@@ -373,12 +355,7 @@ def deserialize_model(cls, api_response: dict[str, Any]) -> dict[str, Any]:
373355
# Appearance
374356
init_data.update(PlotHeight.deserialize(visualize))
375357

376-
# Annotations
377-
init_data["text_annotations"] = TextAnnotation.deserialize_model(
378-
visualize.get("text-annotations")
379-
)
380-
init_data["range_annotations"] = RangeAnnotation.deserialize_model(
381-
visualize.get("range-annotations")
382-
)
358+
# Annotations (from mixin)
359+
init_data.update(cls._deserialize_annotations(visualize))
383360

384361
return init_data

datawrapper/charts/bar.py

Lines changed: 7 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@
55

66
from .base import BaseChart
77
from .enums import DateFormat, NumberFormat, ReplaceFlagsType, ValueLabelAlignment
8-
from .models import RangeAnnotation, TextAnnotation
9-
from .serializers import ColorCategory, CustomRange, ModelListSerializer, ReplaceFlags
8+
from .models import AnnotationsMixin
9+
from .serializers import ColorCategory, CustomRange, ReplaceFlags
1010

1111

1212
class BarOverlay(BaseModel):
@@ -75,7 +75,7 @@ class BarOverlay(BaseModel):
7575
)
7676

7777

78-
class BarChart(BaseChart):
78+
class BarChart(AnnotationsMixin, BaseChart):
7979
"""A base class for the Datawrapper API's bar chart."""
8080

8181
model_config = ConfigDict(
@@ -384,20 +384,6 @@ class BarChart(BaseChart):
384384
description="A list of the highlighted series",
385385
)
386386

387-
#: A list of text annotations to display on the chart
388-
text_annotations: list[TextAnnotation | dict[Any, Any]] = Field(
389-
default_factory=list,
390-
alias="text-annotations",
391-
description="A list of text annotations to display on the chart",
392-
)
393-
394-
#: A list of range annotations to display on the chart
395-
range_annotations: list[RangeAnnotation | dict[Any, Any]] = Field(
396-
default_factory=list,
397-
alias="range-annotations",
398-
description="A list of range annotations to display on the chart",
399-
)
400-
401387
@field_validator("replace_flags")
402388
@classmethod
403389
def validate_replace_flags(
@@ -460,15 +446,12 @@ def serialize_model(self) -> dict:
460446
"overlays": [],
461447
# Annotations
462448
"highlighted-series": self.highlighted_series,
463-
"text-annotations": ModelListSerializer.serialize(
464-
self.text_annotations, TextAnnotation
465-
),
466-
"range-annotations": ModelListSerializer.serialize(
467-
self.range_annotations, RangeAnnotation
468-
),
469449
}
470450
)
471451

452+
# Add annotations
453+
model["metadata"]["visualize"].update(self._serialize_annotations())
454+
472455
# Add the overlays, if any
473456
for overlay_obj in self.overlays:
474457
# If the overlay is a dictionary, validate it and convert it to a BarOverlay object
@@ -608,12 +591,6 @@ def deserialize_model(cls, api_response: dict[str, Any]) -> dict[str, Any]:
608591
if "highlighted-series" in visualize:
609592
init_data["highlighted_series"] = visualize["highlighted-series"]
610593

611-
# Annotations
612-
init_data["text_annotations"] = TextAnnotation.deserialize_model(
613-
visualize.get("text-annotations")
614-
)
615-
init_data["range_annotations"] = RangeAnnotation.deserialize_model(
616-
visualize.get("range-annotations")
617-
)
594+
init_data.update(cls._deserialize_annotations(visualize))
618595

619596
return init_data

datawrapper/charts/column.py

Lines changed: 9 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -14,24 +14,27 @@
1414
ValueLabelPlacement,
1515
)
1616
from .models import (
17+
AnnotationsMixin,
1718
CustomRangeMixin,
1819
CustomTicksMixin,
1920
GridDisplayMixin,
2021
GridFormatMixin,
21-
RangeAnnotation,
22-
TextAnnotation,
2322
)
2423
from .serializers import (
2524
ColorCategory,
26-
ModelListSerializer,
2725
NegativeColor,
2826
PlotHeight,
2927
ValueLabels,
3028
)
3129

3230

3331
class ColumnChart(
34-
GridDisplayMixin, GridFormatMixin, CustomRangeMixin, CustomTicksMixin, BaseChart
32+
AnnotationsMixin,
33+
GridDisplayMixin,
34+
GridFormatMixin,
35+
CustomRangeMixin,
36+
CustomTicksMixin,
37+
BaseChart,
3538
):
3639
"""A base class for the Datawrapper API's column chart."""
3740

@@ -184,24 +187,6 @@ class ColumnChart(
184187
description="Where to place the value labels",
185188
)
186189

187-
#
188-
# Annotations
189-
#
190-
191-
#: A list of text annotations to display on the chart
192-
text_annotations: list[TextAnnotation | dict[Any, Any]] = Field(
193-
default_factory=list,
194-
alias="text-annotations",
195-
description="A list of text annotations to display on the chart",
196-
)
197-
198-
#: A list of range annotations to display on the chart
199-
range_annotations: list[RangeAnnotation | dict[Any, Any]] = Field(
200-
default_factory=list,
201-
alias="range-annotations",
202-
description="A list of range annotations to display on the chart",
203-
)
204-
205190
@field_validator("plot_height_mode")
206191
@classmethod
207192
def validate_plot_height_mode(cls, v: PlotHeightMode | str) -> PlotHeightMode | str:
@@ -362,16 +347,10 @@ def serialize_model(self) -> dict:
362347
placement=self.value_labels_placement,
363348
chart_type="column",
364349
),
365-
# Annotations
366-
"text-annotations": ModelListSerializer.serialize(
367-
self.text_annotations, TextAnnotation
368-
),
369-
"range-annotations": ModelListSerializer.serialize(
370-
self.range_annotations, RangeAnnotation
371-
),
372350
}
373351

374352
model["metadata"]["visualize"].update(visualize_data)
353+
model["metadata"]["visualize"].update(self._serialize_annotations())
375354

376355
# Return the serialized data
377356
return model
@@ -437,11 +416,6 @@ def deserialize_model(cls, api_response: dict[str, Any]) -> dict[str, Any]:
437416
init_data.update(ValueLabels.deserialize(visualize, chart_type="column"))
438417

439418
# Annotations
440-
init_data["text_annotations"] = TextAnnotation.deserialize_model(
441-
visualize.get("text-annotations")
442-
)
443-
init_data["range_annotations"] = RangeAnnotation.deserialize_model(
444-
visualize.get("range-annotations")
445-
)
419+
init_data.update(cls._deserialize_annotations(visualize))
446420

447421
return init_data

datawrapper/charts/line.py

Lines changed: 9 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -25,12 +25,11 @@
2525
SymbolStyle,
2626
)
2727
from .models import (
28+
AnnotationsMixin,
2829
CustomRangeMixin,
2930
CustomTicksMixin,
3031
GridDisplayMixin,
3132
GridFormatMixin,
32-
RangeAnnotation,
33-
TextAnnotation,
3433
)
3534
from .serializers import (
3635
ColorCategory,
@@ -527,7 +526,12 @@ def deserialize_model(cls, line_name: str, line_config: dict) -> dict[str, Any]:
527526

528527

529528
class LineChart(
530-
GridDisplayMixin, GridFormatMixin, CustomRangeMixin, CustomTicksMixin, BaseChart
529+
AnnotationsMixin,
530+
GridDisplayMixin,
531+
GridFormatMixin,
532+
CustomRangeMixin,
533+
CustomTicksMixin,
534+
BaseChart,
531535
):
532536
"""A base class for the Datawrapper API's line chart."""
533537

@@ -758,24 +762,6 @@ def validate_plot_height_mode(cls, v: PlotHeightMode | str) -> PlotHeightMode |
758762
raise ValueError(f"Invalid value: {v}. Must be one of {valid_values}")
759763
return v
760764

761-
#
762-
# Annotations
763-
#
764-
765-
#: A list of text annotations to display on the chart
766-
text_annotations: list[TextAnnotation | dict[Any, Any]] = Field(
767-
default_factory=list,
768-
alias="text-annotations",
769-
description="A list of text annotations to display on the chart",
770-
)
771-
772-
#: A list of range annotations to display on the chart
773-
range_annotations: list[RangeAnnotation | dict[Any, Any]] = Field(
774-
default_factory=list,
775-
alias="range-annotations",
776-
description="A list of range annotations to display on the chart",
777-
)
778-
779765
@model_serializer
780766
def serialize_model(self) -> dict:
781767
"""Serialize the model to a dictionary."""
@@ -817,18 +803,13 @@ def serialize_model(self) -> dict:
817803
),
818804
# Initialize empty structures
819805
"lines": {},
820-
"text-annotations": ModelListSerializer.serialize(
821-
self.text_annotations, TextAnnotation
822-
),
823-
"range-annotations": ModelListSerializer.serialize(
824-
self.range_annotations, RangeAnnotation
825-
),
826806
"custom-area-fills": ModelListSerializer.serialize(
827807
self.area_fills, AreaFill
828808
),
829809
}
830810

831811
model["metadata"]["visualize"].update(visualize_data)
812+
model["metadata"]["visualize"].update(self._serialize_annotations())
832813

833814
# Add line configurations
834815
for line_obj in self.lines:
@@ -933,11 +914,6 @@ def deserialize_model(cls, api_response: dict[str, Any]) -> dict[str, Any]:
933914
init_data.update(PlotHeight.deserialize(visualize))
934915

935916
# Annotations
936-
init_data["text_annotations"] = TextAnnotation.deserialize_model(
937-
visualize.get("text-annotations")
938-
)
939-
init_data["range_annotations"] = RangeAnnotation.deserialize_model(
940-
visualize.get("range-annotations")
941-
)
917+
init_data.update(cls._deserialize_annotations(visualize))
942918

943919
return init_data

datawrapper/charts/models/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
Visualize,
1111
)
1212
from .mixins import (
13+
AnnotationsMixin,
1314
CustomRangeMixin,
1415
CustomTicksMixin,
1516
GridDisplayMixin,
@@ -27,6 +28,7 @@
2728

2829
__all__ = [
2930
"Annotate",
31+
"AnnotationsMixin",
3032
"ColumnFormat",
3133
"ColumnFormatList",
3234
"ConnectorLine",

0 commit comments

Comments
 (0)