Adds overloaded signatures for Cluster.search. Fixes #4

pull/5/head
Christian Keil 2021-11-05 10:30:45 +01:00
parent f9f8a7bdf3
commit a5edb31f2b
1 changed files with 10 additions and 2 deletions

View File

@ -8,7 +8,8 @@ import sys
from collections.abc import Mapping from collections.abc import Mapping
from glob import glob from glob import glob
import re import re
from typing import List, Dict, Optional, Any, Tuple, Iterator from typing import List, Dict, Optional, Any, Tuple, Iterator, overload, Union
from typing_extensions import Literal
try: try:
import jsonschema # type: ignore import jsonschema # type: ignore
@ -225,7 +226,14 @@ class Cluster(Mapping): # type: ignore
raise PyMISPGalaxiesError("Duplicate value ({}) in cluster: {}".format(new_cluster_value.value, self.name)) raise PyMISPGalaxiesError("Duplicate value ({}) in cluster: {}".format(new_cluster_value.value, self.name))
self.cluster_values[new_cluster_value.value] = new_cluster_value self.cluster_values[new_cluster_value.value] = new_cluster_value
def search(self, query: str, return_tags: bool=False) -> List[str]: @overload
def search(self, query: str, return_tags: Literal[False]=False) -> List[ClusterValue]: ...
@overload
def search(self, query: str, return_tags: Literal[True]) -> List[str]: ...
@overload
def search(self, query: str, return_tags: bool) -> Union[List[ClusterValue], List[str]]: ...
def search(self, query: str, return_tags: bool=False) -> Union[List[ClusterValue], List[str]]:
matching = [] matching = []
for v in self.values(): for v in self.values():
if [s for s in v.searchable if query.lower() in s.lower()]: if [s for s in v.searchable if query.lower() in s.lower()]: