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 glob import glob
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:
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))
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 = []
for v in self.values():
if [s for s in v.searchable if query.lower() in s.lower()]: