Skip to content

Commit

Permalink
Spin type in es query op to literal
Browse files Browse the repository at this point in the history
  • Loading branch information
munrojm committed May 11, 2021
1 parent 8d90b54 commit e1d6d86
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions src/mp_api/routes/electronic_structure/query_operators.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Optional
from typing import Optional, Union, Literal
from fastapi import Query, HTTPException
from pymatgen.analysis.magnetism.analyzer import Ordering
from pymatgen.electronic_structure.core import Spin, OrbitalType
Expand Down Expand Up @@ -158,8 +158,9 @@ def query(
projection_type: Optional[DOSProjectionType] = Query(
None, description="Projection type for the density of states data.",
),
spin: Optional[Spin] = Query(
None, description="Spin channel for density of states data.",
spin: Optional[Union[Literal["1", "-1"], Spin]] = Query(
None,
description="Spin channel for density of states data. '1' corresponds to spin up.",
),
element: Optional[Element] = Query(
None, description="Element type for projected density of states data.",
Expand All @@ -186,6 +187,9 @@ def query(

crit = defaultdict(dict) # type: dict

if isinstance(spin, str):
spin = Spin(int(spin))

if projection_type is not None:
if spin is None:
raise HTTPException(
Expand Down

0 comments on commit e1d6d86

Please sign in to comment.