future annotations

Author: janosh

crystal_toolkit/__init__.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 import json
 import os as _os
 from collections import defaultdict
 from pathlib import Path
+from typing import Any
 
 # pleasant hack to support MSONable objects in Dash callbacks natively
 from monty.json import MSONable
@@ -30,7 +33,7 @@ def to_plotly_json(self):
 
 
 # Populate the default values from the JSON file
-_DEFAULTS = defaultdict()
+_DEFAULTS: dict[str, Any] = defaultdict()
 default_js = _os.path.join(
     _os.path.join(_os.path.dirname(_os.path.abspath(__file__))), "./", "defaults.json"
 )

crystal_toolkit/apps/main.py

@@ -1,9 +1,10 @@
+from __future__ import annotations
+
 import logging
 import os
 import warnings
 from random import choice
 from time import time
-from typing import Optional
 from urllib import parse
 from uuid import uuid4
 
@@ -448,7 +449,7 @@ def update_search_term_on_page_load(href: str) -> str:
     [Input(search_component.id("input"), "value")],
     [State(search_component.id("input"), "n_submit")],
 )
-def perform_search_on_page_load(search_term: str, n_submit: Optional[int]):
+def perform_search_on_page_load(search_term: str, n_submit: int | None):
     """
     Loading with an mpid in the URL requires populating the search term with
     the mpid, this callback forces that search to then take place by force updating
@@ -468,7 +469,7 @@ def perform_search_on_page_load(search_term: str, n_submit: Optional[int]):
 
 
 @app.callback(Output("url", "pathname"), [Input(search_component.id(), "data")])
-def update_url_pathname_from_search_term(mpid: Optional[str]) -> str:
+def update_url_pathname_from_search_term(mpid: str | None) -> str:
     """
     Updates the URL from the search term. Technically a circular callback,
     this is done to prevent the app seeming inconsistent from the end user.
@@ -488,7 +489,7 @@ def update_url_pathname_from_search_term(mpid: Optional[str]) -> str:
     Output(transformation_component.id("input_structure"), "data"),
     [Input(search_component.id(), "data"), Input(upload_component.id(), "data")],
 )
-def master_update_structure(search_mpid: Optional[str], upload_data: Optional[str]):
+def master_update_structure(search_mpid: str | None, upload_data: dict | None):
     """
     A new structure is loaded either from the search component or from the
     upload component. This callback triggers the update, and uses the callback

crystal_toolkit/components/bandstructure.py

@@ -26,19 +26,9 @@
     Loading,
     MessageBody,
     MessageContainer,
-    bandstructure_symm_line,
-    bs,
-    bsml,
     dcc,
-    density_of_states,
-    dos,
-    dos_select,
-    elements,
-    get_bandstructure_traces,
     get_data_list,
-    get_dos_traces,
     html,
-    path_convention,
 )
 
 # Author: Jason Munro

crystal_toolkit/components/structure.py

@@ -1,11 +1,13 @@
+from __future__ import annotations
+
 import re
 import warnings
 from base64 import b64encode
 from collections import OrderedDict
 from itertools import chain, combinations_with_replacement
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Dict, Literal, Optional, Tuple, Union
+from typing import Literal
 
 import numpy as np
 from dash import dash_table as dt
@@ -29,7 +31,7 @@
 
 # TODO: make dangling bonds "stubs"? (fixed length)
 
-DEFAULTS = {
+DEFAULTS: dict[str, str | bool] = {
     "color_scheme": "VESTA",
     "bonding_strategy": "CrystalNN",
     "radius_strategy": "uniform",
@@ -85,16 +87,15 @@ class StructureMoleculeComponent(MPComponent):
 
     def __init__(
         self,
-        struct_or_mol: Optional[
-            Union[Structure, StructureGraph, Molecule, MoleculeGraph]
-        ] = None,
+        struct_or_mol: None
+        | (Structure | StructureGraph | Molecule | MoleculeGraph) = None,
         id: str = None,
         className: str = "box",
-        scene_additions: Optional[Scene] = None,
+        scene_additions: Scene | None = None,
         bonding_strategy: str = DEFAULTS["bonding_strategy"],
-        bonding_strategy_kwargs: Optional[dict] = None,
+        bonding_strategy_kwargs: dict | None = None,
         color_scheme: str = DEFAULTS["color_scheme"],
-        color_scale: Optional[str] = None,
+        color_scale: str | None = None,
         radius_strategy: str = DEFAULTS["radius_strategy"],
         unit_cell_choice: str = DEFAULTS["unit_cell_choice"],
         draw_image_atoms: bool = DEFAULTS["draw_image_atoms"],
@@ -103,8 +104,8 @@ def __init__(
         ],
         hide_incomplete_bonds: bool = DEFAULTS["hide_incomplete_bonds"],
         show_compass: bool = DEFAULTS["show_compass"],
-        scene_settings: Optional[Dict] = None,
-        group_by_site_property: Optional[str] = None,
+        scene_settings: dict | None = None,
+        group_by_site_property: str | None = None,
         show_legend: bool = DEFAULTS["show_legend"],
         show_settings: bool = DEFAULTS["show_settings"],
         show_controls: bool = DEFAULTS["show_controls"],
@@ -907,7 +908,7 @@ def layout(self, size: str = "500px") -> html.Div:
 
     @staticmethod
     def _preprocess_structure(
-        struct_or_mol: Union[Structure, StructureGraph, Molecule, MoleculeGraph],
+        struct_or_mol: Structure | StructureGraph | Molecule | MoleculeGraph,
         unit_cell_choice: Literal[
             "input", "primitive", "conventional", "reduced_niggli", "reduced_lll"
         ] = "input",
@@ -931,10 +932,10 @@ def _preprocess_structure(
 
     @staticmethod
     def _preprocess_input_to_graph(
-        input: Union[Structure, StructureGraph, Molecule, MoleculeGraph],
+        input: Structure | StructureGraph | Molecule | MoleculeGraph,
         bonding_strategy: str = DEFAULTS["bonding_strategy"],
-        bonding_strategy_kwargs: Optional[Dict] = None,
-    ) -> Union[StructureGraph, MoleculeGraph]:
+        bonding_strategy_kwargs: dict | None = None,
+    ) -> StructureGraph | MoleculeGraph:
 
         if isinstance(input, Structure):
 
@@ -1007,8 +1008,8 @@ def _preprocess_input_to_graph(
 
     @staticmethod
     def _get_struct_or_mol(
-        graph: Union[StructureGraph, MoleculeGraph, Structure, Molecule]
-    ) -> Union[Structure, Molecule]:
+        graph: StructureGraph | MoleculeGraph | Structure | Molecule,
+    ) -> Structure | Molecule:
         if isinstance(graph, StructureGraph):
             return graph.structure
         elif isinstance(graph, MoleculeGraph):
@@ -1020,7 +1021,7 @@ def _get_struct_or_mol(
 
     @staticmethod
     def get_scene_and_legend(
-        graph: Optional[Union[StructureGraph, MoleculeGraph]],
+        graph: StructureGraph | MoleculeGraph | None,
         color_scheme=DEFAULTS["color_scheme"],
         color_scale=None,
         radius_strategy=DEFAULTS["radius_strategy"],
@@ -1031,7 +1032,7 @@ def get_scene_and_legend(
         scene_additions=None,
         show_compass=DEFAULTS["show_compass"],
         group_by_site_property=None,
-    ) -> Tuple[Scene, Dict[str, str]]:
+    ) -> tuple[Scene, dict[str, str]]:
 
         scene = Scene(name="StructureMoleculeComponentScene")
 

crystal_toolkit/components/transformations/core.py

@@ -1,6 +1,7 @@
+from __future__ import annotations
+
 import traceback
 import warnings
-from typing import List, Optional
 
 import dash
 import dash_daq as daq
@@ -133,7 +134,7 @@ def container_layout(self, state=None, structure=None) -> html.Div:
 
         return container
 
-    def options_layouts(self, state=None, structure=None) -> List[html.Div]:
+    def options_layouts(self, state=None, structure=None) -> list[html.Div]:
         """
         Return a layout to change the transformation options (that is,
         that controls the args and kwargs that will be passed to pymatgen).
@@ -270,8 +271,8 @@ def update_transformation(enabled, states):
 class AllTransformationsComponent(MPComponent):
     def __init__(
         self,
-        transformations: Optional[List[str]] = None,
-        input_structure_component: Optional[MPComponent] = None,
+        transformations: list[str] | None = None,
+        input_structure_component: MPComponent | None = None,
         *args,
         **kwargs,
     ):

crystal_toolkit/core/legend.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import os
 import warnings
 from collections import defaultdict
 from itertools import chain
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any
 
 import numpy as np
 from matplotlib.cm import get_cmap
@@ -41,11 +43,11 @@ class Legend(MSONable):
 
     def __init__(
         self,
-        site_collection: Union[SiteCollection, Site],
+        site_collection: SiteCollection | Site,
         color_scheme: str = "Jmol",
         radius_scheme: str = "uniform",
         cmap: str = "coolwarm",
-        cmap_range: Optional[Tuple[float, float]] = None,
+        cmap_range: tuple[float, float] | None = None,
     ):
         """
         Create a legend for a given SiteCollection to choose how to
@@ -137,7 +139,7 @@ def __init__(
     @staticmethod
     def generate_accessible_color_scheme_on_the_fly(
         site_collection: SiteCollection,
-    ) -> Dict[str, Dict[str, Tuple[int, int, int]]]:
+    ) -> dict[str, dict[str, tuple[int, int, int]]]:
         """
         e.g. for a color scheme more appropriate for people with color blindness
 
@@ -217,7 +219,7 @@ def generate_accessible_color_scheme_on_the_fly(
     @staticmethod
     def generate_categorical_color_scheme_on_the_fly(
         site_collection: SiteCollection, site_prop_types
-    ) -> Dict[str, Dict[str, Tuple[int, int, int]]]:
+    ) -> dict[str, dict[str, tuple[int, int, int]]]:
         """
         e.g. for Wykcoff
 
@@ -257,7 +259,7 @@ def generate_categorical_color_scheme_on_the_fly(
 
         return color_scheme
 
-    def get_color(self, sp: Union[Specie, Element], site: Optional[Site] = None) -> str:
+    def get_color(self, sp: Specie | Element, site: Site | None = None) -> str:
         """
         Get a color to render a specific species. Optionally, you can provide
         a site for context, since ...
@@ -335,9 +337,7 @@ def get_color(self, sp: Union[Specie, Element], site: Optional[Site] = None) ->
 
         return html5_serialize_simple_color(color)
 
-    def get_radius(
-        self, sp: Union[Specie, Element], site: Optional[Site] = None
-    ) -> float:
+    def get_radius(self, sp: Specie | Element, site: Site | None = None) -> float:
 
         # allow manual override by user
         if site and "display_radius" in site.properties:
@@ -380,7 +380,7 @@ def get_radius(
         return radius
 
     @staticmethod
-    def analyze_site_props(site_collection: SiteCollection) -> Dict[str, List[str]]:
+    def analyze_site_props(site_collection: SiteCollection) -> dict[str, list[str]]:
         """
         Returns: A dictionary with keys "scalar", "matrix", "vector", "categorical"
         and values of a list of site property names corresponding to each type
@@ -400,7 +400,7 @@ def analyze_site_props(site_collection: SiteCollection) -> Dict[str, List[str]]:
         return dict(site_prop_names)
 
     @staticmethod
-    def get_species_str(sp: Union[Specie, Element]) -> str:
+    def get_species_str(sp: Specie | Element) -> str:
         """
         Args:
             sp: Specie or Element
@@ -411,7 +411,7 @@ def get_species_str(sp: Union[Specie, Element]) -> str:
         # and then move this to pymatgen string utils ...
         return unicodeify_species(str(sp))
 
-    def get_legend(self) -> Dict[str, Any]:
+    def get_legend(self) -> dict[str, Any]:
 
         # decide what we want the labels to be
         if self.color_scheme in ("Jmol", "VESTA", "accessible"):