Merge pull request #2398 from dopplershift/stationlookup

Make StationLookup a proper Mapping

src/metpy/io/__init__.py

@@ -1,11 +1,15 @@
 # Copyright (c) 2015,2016,2018,2021 MetPy Developers.
 # Distributed under the terms of the BSD 3-Clause License.
 # SPDX-License-Identifier: BSD-3-Clause
-"""Classes for reading various file formats.
+"""Tools for reading various file formats.
 
-These classes are written to take both file names (for local files) or file-like objects;
-this allows reading files that are already in memory (using :class:`python:io.StringIO`)
-or remote files (using :func:`~python:urllib.request.urlopen`).
+Classes supporting formats are written to take both file names (for local files) or file-like
+objects; this allows reading files that are already in memory
+(using :class:`python:io.StringIO`) or remote files
+(using :func:`~python:urllib.request.urlopen`).
+
+`station_info` is an instance of `StationLookup` to find information about station locations
+(e.g. latitude, longitude, altitude) from various sources.
 """
 
 from .gempak import *  # noqa: F403

src/metpy/io/station_data.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: BSD-3-Clause
 """Pull out station metadata."""
 from collections import ChainMap, namedtuple
+from collections.abc import Mapping
 from functools import cached_property
 
 import numpy as np
@@ -120,8 +121,18 @@ def _read_airports_file(input_file=None):
                          }).to_dict()
 
 
-class StationLookup:
-    """Look up station information from multiple sources."""
+@exporter.export
+class StationLookup(Mapping):
+    """Look up station information from multiple sources.
+
+    This class follows the `Mapping` protocol with station ID as the key. This makes it
+    possible to e.g. iterate over all locations and get all of a certain criteria:
+
+    >>> import metpy.io
+    >>> conus_stations = [s for s in metpy.io.station_info if s.startswith('K')]
+    >>> conus_stations[:3]
+    ['KEET', 'K8A0', 'KALX']
+    """
 
     @cached_property
     def tables(self):
@@ -131,6 +142,14 @@ def tables(self):
                         dict(_read_station_text_file()),
                         dict(_read_airports_file()))
 
+    def __len__(self):
+        """Get the number of stations."""
+        return len(self.tables)
+
+    def __iter__(self):
+        """Allow iteration over the stations."""
+        return iter(self.tables)
+
     def __getitem__(self, stid):
         """Lookup station information from the ID."""
         try:

tests/io/test_station_data.py

@@ -7,7 +7,7 @@
 import pandas as pd
 import pytest
 
-from metpy.io import add_station_lat_lon
+from metpy.io import add_station_lat_lon, station_info
 
 
 def test_add_lat_lon_station_data():
@@ -39,3 +39,19 @@ def test_add_lat_lon_station_data_not_found():
 
     with pytest.raises(KeyError):
         add_station_lat_lon(df)
+
+
+def test_station_lookup_get_station():
+    """Test that you can get a station by ID from the lookup."""
+    assert station_info['KOUN'].id == 'KOUN'
+
+
+def test_station_lookup_len():
+    """Test that you can get the length of the station data."""
+    assert len(station_info) == 13798
+
+
+def test_station_lookup_iter():
+    """Test iterating over the station data."""
+    for stid in station_info:
+        assert stid in station_info