add annotations

src/standard_names/error.py

@@ -1,4 +1,4 @@
-#! /usr/bin/env python
+from collections.abc import Iterable
 
 
 class Error(Exception):
@@ -12,29 +12,29 @@ class BadNameError(Error):
 
     """Error to indicate a poorly-formed standard name."""
 
-    def __init__(self, name):
+    def __init__(self, name: str):
         super().__init__()
         self._name = name
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self._name
 
     @property
-    def name(self):
+    def name(self) -> str:
         return self._name
 
 
 class BadRegistryError(Error):
 
     """Error to indicate a bad NamesRegistry."""
 
-    def __init__(self, names):
+    def __init__(self, names: Iterable[str]):
         super().__init__()
         self._names = tuple(sorted(names))
 
-    def __str__(self):
+    def __str__(self) -> str:
         return "Registry contains invalid names"
 
     @property
-    def names(self):
+    def names(self) -> tuple[str, ...]:
         return self._names

src/standard_names/registry.py

@@ -1,13 +1,19 @@
+from __future__ import annotations
+
 import os
 import warnings
+from collections.abc import Iterable
 from glob import glob
 
+from packaging.version import InvalidVersion
+from packaging.version import Version
+
 from standard_names.error import BadNameError
 from standard_names.error import BadRegistryError
 from standard_names.standardname import StandardName
 
 
-def load_names_from_txt(file_like, onerror="raise"):
+def load_names_from_txt(file_like, onerror: str = "raise") -> set[StandardName]:
     """Load names from a text file.
 
     Parameters
@@ -60,17 +66,16 @@ def load_names_from_txt(file_like, onerror="raise"):
     return names
 
 
-def _strict_version_or_raise(version_str):
-    from packaging.version import InvalidVersion
-    from packaging.version import Version
-
+def _strict_version_or_raise(version_str: str) -> Version:
     try:
         return Version(version_str)
     except InvalidVersion as error:
         raise ValueError(f"{version_str}: Not a version string") from error
 
 
-def _get_latest_names_file(path=None, prefix="names-", suffix=".txt"):
+def _get_latest_names_file(
+    path: str | None = None, prefix: str = "names-", suffix: str = ".txt"
+) -> tuple[str | None, str | None]:
     """Get the most recent version of a names file.
 
     Parameters
@@ -196,17 +201,17 @@ class NamesRegistry:
     >>> sorted(registry.names_with('temperature'))
     ['air__temperature', 'water__temperature']
     >>> registry.names_with(['temperature', 'air'])
-    ['air__temperature']
+    {'air__temperature'}
 
     Use ``match`` to match names using a glob-style pattern.
 
     >>> registry.match('air*')
-    ['air__temperature']
+    {'air__temperature'}
 
     Use ``search`` to do a fuzzy search of the list.
 
     >>> registry.search('air__temp')
-    ['air__temperature']
+    {'air__temperature'}
     """
 
     def __init__(self, *args, **kwds):
@@ -237,12 +242,12 @@ def __init__(self, *args, **kwds):
             else:
                 self._load(path)
 
-    def _load(self, file_like, onerror="raise"):
+    def _load(self, file_like, onerror: str = "raise") -> None:
         for name in load_names_from_txt(file_like, onerror=onerror):
             self.add(name)
 
     @property
-    def version(self):
+    def version(self) -> str:
         """The version of the names database.
 
         Returns
@@ -253,7 +258,7 @@ def version(self):
         return self._version
 
     @property
-    def names(self):
+    def names(self) -> tuple[str, ...]:
         """All names in the registry.
 
         Returns
@@ -264,7 +269,7 @@ def names(self):
         return tuple(self._names)
 
     @property
-    def objects(self):
+    def objects(self) -> tuple[str, ...]:
         """All objects in the registry.
 
         Returns
@@ -275,7 +280,7 @@ def objects(self):
         return tuple(self._objects)
 
     @property
-    def quantities(self):
+    def quantities(self) -> tuple[str, ...]:
         """All quantities in the registry.
 
         Returns
@@ -286,7 +291,7 @@ def quantities(self):
         return tuple(self._quantities)
 
     @property
-    def operators(self):
+    def operators(self) -> tuple[str, ...]:
         """All operators in the registry.
 
         Returns
@@ -297,7 +302,7 @@ def operators(self):
         return tuple(self._operators)
 
     @classmethod
-    def from_path(cls, path):
+    def from_path(cls, path: str) -> NamesRegistry:
         """Create a new registry from a text file.
 
         Parameters
@@ -312,15 +317,16 @@ def from_path(cls, path):
         """
         return cls(path)
 
-    def add(self, name):
+    def add(self, name: str | StandardName) -> None:
         """Add a name to the registry.
 
         Parameters
         ----------
         name : str
             A Standard Name.
         """
-        if not isinstance(name, StandardName):
+        # if not isinstance(name, StandardName):
+        if isinstance(name, str):
             name = StandardName(name)
 
         self._names.add(name.name)
@@ -329,18 +335,18 @@ def add(self, name):
         for op in name.operators:
             self._operators.add(op)
 
-    def __contains__(self, name):
+    def __contains__(self, name: str) -> bool:
         if isinstance(name, StandardName):
             name = name.name
         return name in self._names
 
-    def __len__(self):
+    def __len__(self) -> int:
         return len(self._names)
 
     def __iter__(self):
         yield from self._names
 
-    def search(self, name):
+    def search(self, name: str) -> set[str]:
         """Search the registry for a name.
 
         Parameters
@@ -355,9 +361,9 @@ def search(self, name):
         """
         from difflib import get_close_matches
 
-        return get_close_matches(name, self._names)
+        return set(get_close_matches(name, self._names))
 
-    def match(self, pattern):
+    def match(self, pattern: str) -> set[str]:
         """Search the registry for names that match a pattern.
 
         Parameters
@@ -374,13 +380,9 @@ def match(self, pattern):
         import re
 
         p = re.compile(fnmatch.translate(pattern))
-        names = []
-        for name in self._names:
-            if p.match(name):
-                names.append(name)
-        return names
+        return {name for name in self._names if p.match(name)}
 
-    def names_with(self, parts):
+    def names_with(self, parts: str | Iterable[str]) -> set[str]:
         """Search the registry for names containing words.
 
         Parameters
@@ -396,14 +398,8 @@ def names_with(self, parts):
         if isinstance(parts, str):
             parts = (parts,)
 
-        remaining_names = self._names
-        for part in parts:
-            names = []
-            for name in remaining_names:
-                if part in name:
-                    names.append(name)
-            remaining_names = names
-        return names
+        return {name for name in self._names if all(part in name for part in parts)}
+
 
 
 REGISTRY = NamesRegistry()

src/standard_names/standardname.py

@@ -10,7 +10,7 @@
 #    '^[a-z][a-z0-9_]*[a-z0-9](__)[a-z0-9][a-z0-9_]*[a-z0-9]$'
 
 
-def is_valid_name(name):
+def is_valid_name(name: str) -> bool:
     """Check if a string is a valid standard name.
 
     Parameters
@@ -59,7 +59,7 @@ class StandardName:
 
     re = _PREFIX_REGEX + "(__)" + _SUFFIX_REGEX
 
-    def __init__(self, name):
+    def __init__(self, name: str):
         """Create a standard name object from a string.
 
         Parameters
@@ -77,7 +77,7 @@ def __init__(self, name):
         )
 
     @staticmethod
-    def decompose_name(name):
+    def decompose_name(name: str) -> tuple[str, str, tuple[str, ...]]:
         """Decompose a name into its parts.
 
         Decompose the *name* standard name string into it's constituent
@@ -117,7 +117,9 @@ def decompose_name(name):
         return object_part, quantity_part, operators
 
     @staticmethod
-    def compose_name(object, quantity, operators=()):
+    def compose_name(
+        object: str, quantity: str, operators: tuple[()] | tuple[str, ...] = ()
+    ) -> str:
         """Create a string from the parts of StandardName.
 
         Parameters
@@ -151,7 +153,7 @@ def compose_name(object, quantity, operators=()):
         return "__".join((object, quantity))
 
     @staticmethod
-    def decompose_quantity(quantity_clause):
+    def decompose_quantity(quantity_clause: str) -> tuple[tuple[str, ...], str]:
         """Decompose a quantity into operators and quantities.
 
         Decompose the *quantity_clause* string into operator and base
@@ -177,67 +179,67 @@ def decompose_quantity(quantity_clause):
         return operators, quantity
 
     @property
-    def name(self):
+    def name(self) -> str:
         """The full standard name as a string."""
         return self._name
 
     @property
-    def object(self):
+    def object(self) -> str:
         """The object part of the standard name."""
         return self._object
 
     @object.setter
-    def object(self, value):
+    def object(self, value: str):
         self._object = value
         self._name = StandardName.compose_name(
             self.object, self.quantity, self.operators
         )
 
     @property
-    def quantity(self):
+    def quantity(self) -> str:
         """The quantity part of the standard name."""
         return self._quantity
 
     @quantity.setter
-    def quantity(self, value):
+    def quantity(self, value: str):
         self._quantity = value
         self._name = StandardName.compose_name(
             self.object, self.quantity, self.operators
         )
 
     @property
-    def operators(self):
+    def operators(self) -> tuple[str, ...]:
         """The operator part of the standard name."""
         return self._operators
 
     @operators.setter
-    def operators(self, value):
+    def operators(self, value: str | tuple[str, ...]):
         if isinstance(value, str):
             value = (value,)
         self._operators = value
         self._name = StandardName.compose_name(
             self.object, self.quantity, self.operators
         )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "StandardName(%r)" % self.name
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.name
 
-    def __eq__(self, that):
+    def __eq__(self, that) -> bool:
         return self.name == str(that)
 
-    def __ne__(self, that):
+    def __ne__(self, that) -> bool:
         return self.name != str(that)
 
-    def __lt__(self, that):
+    def __lt__(self, that) -> bool:
         return self.name < str(that)
 
-    def __gt__(self, that):
+    def __gt__(self, that) -> bool:
         return self.name > str(that)
 
-    def __cmp__(self, that):
+    def __cmp__(self, that) -> bool:
         return self.name == str(that)
 
     def __hash__(self):