Skip to content

Commit

Permalink
Tidy optional type hints
Browse files Browse the repository at this point in the history
  • Loading branch information
ElliottKasoar committed Jun 11, 2024
1 parent 9fa4bef commit bc503bb
Showing 1 changed file with 44 additions and 44 deletions.
88 changes: 44 additions & 44 deletions abcd/backends/atoms_opensearch.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from collections.abc import Generator
from datetime import datetime
from typing import Union, Iterable
from typing import Iterable, Optional, Union
import logging
from os import linesep
from pathlib import Path
Expand Down Expand Up @@ -108,29 +108,29 @@ class AtomsModel(AbstractModel):
Attributes
----------
_client: Union[OpenSearch, None]
_client: Optional[OpenSearch]
OpenSearch client.
_index_name: Union[str, None]
_index_name: Optional[str]
OpenSearch index name.
"""

def __init__(
self,
client: Union[OpenSearch, None] = None,
index_name: Union[str, None] = None,
dict: Union[dict, None] = None,
client: Optional[OpenSearch] = None,
index_name: Optional[str] = None,
dict: Optional[dict] = None,
):
"""
Initialises class.
Parameters
----------
client: Union[OpenSearch, None]
OpenSearch client.
index_name: Union[str, None]
OpenSearch index name.
dict: dict
Dictionary of atoms data.
client: Optional[OpenSearch]
OpenSearch client. Default is `None`.
index_name: Optional[str]
OpenSearch index name. Default is `None`.
dict: Optional[dict]
Dictionary of atoms data. Default is `None`.
"""
super().__init__(dict)

Expand All @@ -143,7 +143,7 @@ def from_atoms(
client: OpenSearch,
index_name: str,
atoms: Atoms,
extra_info: Union[dict, None] = None,
extra_info: Optional[dict] = None,
store_calc: bool = True,
) -> AtomsModel:
"""
Expand All @@ -157,7 +157,7 @@ def from_atoms(
OpenSearch index name.
atoms: Atoms
Atoms data to be stored.
extra_info: Union[dict, None], optional
extra_info: Optional[dict]
Extra information to store in the document with the atoms data.
Default is `None`.
store_calc: bool, optional
Expand Down Expand Up @@ -316,13 +316,13 @@ def info(self):
"type": "opensearch",
}

def delete(self, query: Union[dict, str, None] = None):
def delete(self, query: Optional[dict, str] = None):
"""
Deletes documents from the database.
Parameters
----------
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to be deleted. Default is `None`.
"""
query = self.parser(query)
Expand Down Expand Up @@ -375,7 +375,7 @@ def save_bulk(self, actions: Iterable, **kwargs):
def push(
self,
atoms: Union[Atoms, Iterable],
extra_info: Union[dict, str, list, None] = None,
extra_info: Optional[dict, str, list] = None,
store_calc: bool = True,
**kwargs,
):
Expand All @@ -385,7 +385,7 @@ def push(
Parameters
----------
atoms: Union[Atoms, Iterable]
extra_info: Union[dict, str, None], optional
extra_info: Optional[dict, str], optional
Extra information to store in the document with the atoms data.
Default is `None`.
store_calc: bool, optional
Expand Down Expand Up @@ -430,7 +430,7 @@ def push(
def upload(
self,
file: Path,
extra_infos: Union[Iterable, dict, None] = None,
extra_infos: Optional[Iterable, dict] = None,
store_calc: bool = True,
):
"""
Expand All @@ -440,7 +440,7 @@ def upload(
----------
file: Path
Path to file to be uploaded
extra_infos: Union[Iterable, dict, None], optional
extra_infos: Optional[Iterable, dict]
Extra information to store in the document with the atoms data.
Default is `None`.
store_calc: bool, optional
Expand All @@ -462,14 +462,14 @@ def upload(
self.push(data, extra_info, store_calc=store_calc)

def get_items(
self, query: Union[dict, str, None] = None
self, query: Optional[dict, str] = None
) -> Generator[dict, None, None]:
"""
Get data as a dictionary from documents in the database.
Parameters
----------
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to get data from. Default is `None`.
Returns
Expand All @@ -489,14 +489,14 @@ def get_items(
yield {"_id": hit["_id"], **hit["_source"]}

def get_atoms(
self, query: Union[dict, str, None] = None
self, query: Optional[dict, str] = None
) -> Generator[Atoms, None, None]:
"""
Get data as Atoms object from documents in the database.
Parameters
----------
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to get data from. Default is `None`.
Returns
Expand All @@ -515,13 +515,13 @@ def get_atoms(
):
yield AtomsModel(None, None, hit["_source"]).to_ase()

def count(self, query: Union[dict, str, None] = None, timeout=30.0) -> int:
def count(self, query: Optional[dict, str] = None, timeout=30.0) -> int:
"""
Counts number of documents in the database.
Parameters
----------
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to be counted. Default is `None`.
timeout: float
Timeout for request in seconds.
Expand All @@ -543,7 +543,7 @@ def count(self, query: Union[dict, str, None] = None, timeout=30.0) -> int:
def _get_props_from_source(
self,
names: Union[str, list[str]],
query: Union[dict, str, None] = None,
query: Optional[dict, str] = None,
) -> dict:
"""
Gets all values of specified properties using the original data from _source.
Expand All @@ -552,7 +552,7 @@ def _get_props_from_source(
----------
names: Union[str, list[str]]
Name or list of names of properties to return.
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to get properties from. Default is `None`.
Returns
Expand Down Expand Up @@ -580,7 +580,7 @@ def property(
self,
names: Union[str, list[str]],
allow_flatten: bool = True,
query: Union[dict, str, None] = None,
query: Optional[dict, str] = None,
) -> Union[dict, list]:
"""
Gets all values of specified properties for matching documents in the database.
Expand All @@ -592,7 +592,7 @@ def property(
allow_flatten: bool = True
Whether to allow arrays to be returned flattened. There is no guarantee
for the order of returned values. Default is `True`.
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to get properties from. Default is `None`.
Returns
Expand Down Expand Up @@ -644,15 +644,15 @@ def property(
return props[names[0]]
return props

def count_property(self, name, query: Union[dict, str, None] = None) -> dict:
def count_property(self, name, query: Optional[dict, str] = None) -> dict:
"""
Counts values of a specified property for matching documents in the
database. This method much faster than performing a Count on the list
returned by self.property, so this method should be used preferentially.
Parameters
----------
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to count properties from. Default is `None`.
Returns
Expand Down Expand Up @@ -684,14 +684,14 @@ def count_property(self, name, query: Union[dict, str, None] = None) -> dict:

return prop

def properties(self, query: Union[dict, str, None] = None) -> dict:
def properties(self, query: Optional[dict, str] = None) -> dict:
"""
Gets lists of all properties from matching documents, separated into
info, derived, and array properties.
Parameters
----------
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to get properties from. Default is `None`.
Returns
Expand Down Expand Up @@ -777,13 +777,13 @@ def get_type_of_property(self, prop: str, category: str) -> str:
return "vector({})".format(map_types[type(data[0])])
return "scalar({})".format(map_types[type(data)])

def count_properties(self, query: Union[dict, str, None] = None) -> dict:
def count_properties(self, query: Optional[dict, str] = None) -> dict:
"""
Counts all properties from matching documents.
Parameters
----------
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to count properties from. Default is `None`.
Returns
Expand Down Expand Up @@ -836,15 +836,15 @@ def count_properties(self, query: Union[dict, str, None] = None) -> dict:

return properties

def add_property(self, data: dict, query: Union[dict, str, None] = None):
def add_property(self, data: dict, query: Optional[dict, str] = None):
"""
Adds properties to matching documents.
Parameters
----------
data: dict
Property key-value pairs to be added to matching documents.
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to add properties to. Default is `None`.
"""
query = self.parser(query)
Expand All @@ -869,7 +869,7 @@ def add_property(self, data: dict, query: Union[dict, str, None] = None):
)

def rename_property(
self, name: str, new_name: str, query: Union[dict, str, None] = None
self, name: str, new_name: str, query: Optional[dict, str] = None
):
"""
Renames property for all matching documents.
Expand All @@ -880,7 +880,7 @@ def rename_property(
Current name of property to be renamed.
new_name: str
New name of property to be renamed.
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to rename property. Default is `None`.
"""
query = self.parser(query)
Expand All @@ -906,15 +906,15 @@ def rename_property(

self.client.update_by_query(index=self.index_name, body=body)

def delete_property(self, name: str, query: Union[dict, str, None] = None):
def delete_property(self, name: str, query: Optional[dict, str] = None):
"""
Deletes property from all matching documents.
Parameters
----------
name: str
Name of property to be deleted from documents.
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents to have property deleted. Default is `None`.
"""
query = self.parser(query)
Expand All @@ -940,7 +940,7 @@ def delete_property(self, name: str, query: Union[dict, str, None] = None):
self.client.update_by_query(index=self.index_name, body=body)

def hist(
self, name: str, query: Union[dict, str, None] = None, **kwargs
self, name: str, query: Optional[dict, str] = None, **kwargs
) -> Union[dict, None]:
"""
Calculate histogram statistics for a property from all matching documents.
Expand All @@ -949,7 +949,7 @@ def hist(
----------
name: str
Name of property.
query: Union[dict, str, None]
query: Optional[dict, str]
Query to filter documents. Default is `None`.
Returns
Expand Down

0 comments on commit bc503bb

Please sign in to comment.