mirror of
https://github.com/OpenBB-finance/OpenBB.git
synced 2026-10-30 14:02:06 +00:00
[Feature] Custom choices (#6169)
* custom choices * remove unnecessary choices from the extra info * static assets * simpler way to pop choices on extra * Getting json_schema_extra without changing the original dict --------- Co-authored-by: montezdesousa <79287829+montezdesousa@users.noreply.github.com>
This commit is contained in:
@@ -2,6 +2,8 @@ import sys
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict, Optional
|
||||
|
||||
from typing_extensions import LiteralString
|
||||
|
||||
# `slots` is available on Python >= 3.10
|
||||
if sys.version_info >= (3, 10):
|
||||
slots_true = {"slots": True}
|
||||
@@ -24,3 +26,10 @@ class OpenBBCustomParameter(BaseMetadata):
|
||||
"""Custom parameter for OpenBB."""
|
||||
|
||||
description: Optional[str] = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, **slots_true)
|
||||
class OpenBBCustomChoices(BaseMetadata):
|
||||
"""Custom choices for OpenBB."""
|
||||
|
||||
choices: Optional[LiteralString] = None
|
||||
|
||||
@@ -36,7 +36,10 @@ from starlette.routing import BaseRoute
|
||||
from typing_extensions import Annotated, _AnnotatedAlias
|
||||
|
||||
from openbb_core.app.extension_loader import ExtensionLoader, OpenBBGroups
|
||||
from openbb_core.app.model.custom_parameter import OpenBBCustomParameter
|
||||
from openbb_core.app.model.custom_parameter import (
|
||||
OpenBBCustomChoices,
|
||||
OpenBBCustomParameter,
|
||||
)
|
||||
from openbb_core.app.provider_interface import ProviderInterface
|
||||
from openbb_core.app.router import CommandMap, RouterLoader
|
||||
from openbb_core.app.static.utils.console import Console
|
||||
@@ -330,9 +333,7 @@ class ImportDefinition:
|
||||
hint_type_list = cls.get_path_hint_type_list(path=path)
|
||||
code = "from openbb_core.app.static.container import Container"
|
||||
code += "\nfrom openbb_core.app.model.obbject import OBBject"
|
||||
code += (
|
||||
"\nfrom openbb_core.app.model.custom_parameter import OpenBBCustomParameter"
|
||||
)
|
||||
code += "\nfrom openbb_core.app.model.custom_parameter import OpenBBCustomParameter, OpenBBCustomChoices"
|
||||
|
||||
# These imports were not detected before build, so we add them manually and
|
||||
# ruff --fix the resulting code to remove unused imports.
|
||||
@@ -500,7 +501,10 @@ class MethodDefinition:
|
||||
"""Get json schema extra."""
|
||||
field_default = getattr(field, "default", None)
|
||||
if field_default:
|
||||
return getattr(field_default, "json_schema_extra", {})
|
||||
# Getting json_schema_extra without changing the original dict
|
||||
json_schema_extra = getattr(field_default, "json_schema_extra", {}).copy()
|
||||
json_schema_extra.pop("choices", None)
|
||||
return json_schema_extra
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
@@ -604,10 +608,10 @@ class MethodDefinition:
|
||||
return MethodDefinition.reorder_params(params=formatted)
|
||||
|
||||
@staticmethod
|
||||
def add_field_descriptions(
|
||||
def add_field_custom_annotations(
|
||||
od: OrderedDict[str, Parameter], model_name: Optional[str] = None
|
||||
):
|
||||
"""Add the field description to the param signature."""
|
||||
"""Add the field custom description and choices to the param signature as annotations."""
|
||||
if model_name:
|
||||
available_fields: Dict[str, Field] = (
|
||||
ProviderInterface().params[model_name]["standard"].__dataclass_fields__
|
||||
@@ -617,16 +621,28 @@ class MethodDefinition:
|
||||
if param not in available_fields:
|
||||
continue
|
||||
|
||||
field = available_fields[param]
|
||||
field_default = available_fields[param].default
|
||||
|
||||
new_value = value.replace(
|
||||
annotation=Annotated[
|
||||
value.annotation,
|
||||
OpenBBCustomParameter(
|
||||
description=getattr(field.default, "description", "")
|
||||
),
|
||||
],
|
||||
choices = getattr(field_default, "json_schema_extra", {}).get(
|
||||
"choices", []
|
||||
)
|
||||
description = getattr(field_default, "description", "")
|
||||
|
||||
if choices:
|
||||
new_value = value.replace(
|
||||
annotation=Annotated[
|
||||
value.annotation,
|
||||
OpenBBCustomParameter(description=description),
|
||||
OpenBBCustomChoices(choices=choices),
|
||||
],
|
||||
)
|
||||
else:
|
||||
new_value = value.replace(
|
||||
annotation=Annotated[
|
||||
value.annotation,
|
||||
OpenBBCustomParameter(description=description),
|
||||
],
|
||||
)
|
||||
|
||||
od[param] = new_value
|
||||
|
||||
@@ -667,7 +683,7 @@ class MethodDefinition:
|
||||
model_name: Optional[str] = None,
|
||||
) -> str:
|
||||
"""Build the command method signature."""
|
||||
MethodDefinition.add_field_descriptions(
|
||||
MethodDefinition.add_field_custom_annotations(
|
||||
od=formatted_params, model_name=model_name
|
||||
) # this modified `od` in place
|
||||
func_params = MethodDefinition.build_func_params(formatted_params)
|
||||
|
||||
@@ -74,7 +74,10 @@ CPI_FREQUENCY = Literal["monthly", "quarter", "annual"]
|
||||
class ConsumerPriceIndexQueryParams(QueryParams):
|
||||
"""CPI Query."""
|
||||
|
||||
country: str = Field(description=QUERY_DESCRIPTIONS.get("country"))
|
||||
country: str = Field(
|
||||
description=QUERY_DESCRIPTIONS.get("country"),
|
||||
choices=CPI_COUNTRIES, # type: ignore
|
||||
)
|
||||
units: CPI_UNITS = Field(
|
||||
default="growth_same",
|
||||
description=QUERY_DESCRIPTIONS.get("units", "")
|
||||
|
||||
@@ -3,7 +3,10 @@
|
||||
import datetime
|
||||
from typing import List, Literal, Optional, Union
|
||||
|
||||
from openbb_core.app.model.custom_parameter import OpenBBCustomParameter
|
||||
from openbb_core.app.model.custom_parameter import (
|
||||
OpenBBCustomChoices,
|
||||
OpenBBCustomParameter,
|
||||
)
|
||||
from openbb_core.app.model.obbject import OBBject
|
||||
from openbb_core.app.static.container import Container
|
||||
from openbb_core.app.static.utils.decorators import exception_handler, validate
|
||||
@@ -240,6 +243,59 @@ class ROUTER_economy(Container):
|
||||
OpenBBCustomParameter(
|
||||
description="The country to get data. Multiple items allowed for provider(s): fred."
|
||||
),
|
||||
OpenBBCustomChoices(
|
||||
choices=[
|
||||
"australia",
|
||||
"austria",
|
||||
"belgium",
|
||||
"brazil",
|
||||
"bulgaria",
|
||||
"canada",
|
||||
"chile",
|
||||
"china",
|
||||
"croatia",
|
||||
"cyprus",
|
||||
"czech_republic",
|
||||
"denmark",
|
||||
"estonia",
|
||||
"euro_area",
|
||||
"finland",
|
||||
"france",
|
||||
"germany",
|
||||
"greece",
|
||||
"hungary",
|
||||
"iceland",
|
||||
"india",
|
||||
"indonesia",
|
||||
"ireland",
|
||||
"israel",
|
||||
"italy",
|
||||
"japan",
|
||||
"korea",
|
||||
"latvia",
|
||||
"lithuania",
|
||||
"luxembourg",
|
||||
"malta",
|
||||
"mexico",
|
||||
"netherlands",
|
||||
"new_zealand",
|
||||
"norway",
|
||||
"poland",
|
||||
"portugal",
|
||||
"romania",
|
||||
"russian_federation",
|
||||
"slovak_republic",
|
||||
"slovakia",
|
||||
"slovenia",
|
||||
"south_africa",
|
||||
"spain",
|
||||
"sweden",
|
||||
"switzerland",
|
||||
"turkey",
|
||||
"united_kingdom",
|
||||
"united_states",
|
||||
]
|
||||
),
|
||||
],
|
||||
units: Annotated[
|
||||
Literal["growth_previous", "growth_same", "index_2015"],
|
||||
|
||||
Reference in New Issue
Block a user