diff --git a/Lib/selectors.py b/Lib/selectors.py --- a/Lib/selectors.py +++ b/Lib/selectors.py @@ -6,7 +6,7 @@ This module allows high-level and effici from abc import ABCMeta, abstractmethod -from collections import namedtuple +from collections import namedtuple, Mapping import functools import select import sys @@ -44,6 +44,24 @@ SelectorKey = namedtuple('SelectorKey', selected event mask and attached data.""" +class _SelectorMapping(Mapping): + + def __init__(self, selector): + self._fd_to_key = selector._fd_to_key + + def __len__(self): + return len(self._fd_to_key) + + def __getitem__(self, fileobj): + try: + return self._fd_to_key[_fileobj_to_fd(fileobj)] + except KeyError: + raise KeyError("{!r} is not registered".format(fileobj)) from None + + def __iter__(self): + return iter(self._fd_to_key) + + class BaseSelector(metaclass=ABCMeta): """Base selector class. @@ -151,16 +169,9 @@ class BaseSelector(metaclass=ABCMeta): """ self._fd_to_key.clear() - def get_key(self, fileobj): - """Return the key associated to a registered file object. - - Returns: - SelectorKey for this file object - """ - try: - return self._fd_to_key[_fileobj_to_fd(fileobj)] - except KeyError: - raise KeyError("{!r} is not registered".format(fileobj)) from None + def get_map(self): + """Get a mapping of file descriptors to selector keys.""" + return _SelectorMapping(self) def __enter__(self): return self diff --git a/Lib/test/test_selectors.py b/Lib/test/test_selectors.py --- a/Lib/test/test_selectors.py +++ b/Lib/test/test_selectors.py @@ -106,7 +106,7 @@ class BaseSelectorTestCase(unittest.Test # modify events key2 = s.modify(rd, selectors.EVENT_WRITE) self.assertNotEqual(key.events, key2.events) - self.assertEqual(key2, s.get_key(rd)) + self.assertEqual(key2, s.get_map()[rd]) s.unregister(rd) @@ -118,7 +118,7 @@ class BaseSelectorTestCase(unittest.Test key2 = s.modify(rd, selectors.EVENT_READ, d2) self.assertEqual(key.events, key2.events) self.assertNotEqual(key.data, key2.data) - self.assertEqual(key2, s.get_key(rd)) + self.assertEqual(key2, s.get_map()[rd]) self.assertEqual(key2.data, d2) # modify unknown file obj @@ -136,10 +136,12 @@ class BaseSelectorTestCase(unittest.Test s.register(wr, selectors.EVENT_WRITE) s.close() - self.assertRaises(KeyError, s.get_key, rd) - self.assertRaises(KeyError, s.get_key, wr) + with self.assertRaises(KeyError): + s.get_map()[rd] + with self.assertRaises(KeyError): + s.get_map()[wr] - def test_get_key(self): + def test_get_map(self): s = self.SELECTOR() self.addCleanup(s.close) @@ -147,11 +149,23 @@ class BaseSelectorTestCase(unittest.Test self.addCleanup(rd.close) self.addCleanup(wr.close) + keys = s.get_map() + self.assertEqual(len(keys), 0) + self.assertEqual(list(keys), []) key = s.register(rd, selectors.EVENT_READ, "data") - self.assertEqual(key, s.get_key(rd)) + self.assertIn(rd, keys) + self.assertEqual(key, keys[rd]) + self.assertEqual(len(keys), 1) + self.assertEqual(list(keys), [rd.fileno()]) + self.assertEqual(list(keys.values()), [key]) # unknown file obj - self.assertRaises(KeyError, s.get_key, 999999) + with self.assertRaises(KeyError): + keys[999999] + + # Read-only mapping + with self.assertRaises(TypeError): + del keys[rd] def test_select(self): s = self.SELECTOR() @@ -185,8 +199,8 @@ class BaseSelectorTestCase(unittest.Test sel.register(rd, selectors.EVENT_READ) sel.register(wr, selectors.EVENT_WRITE) - self.assertRaises(KeyError, s.get_key, rd) - self.assertRaises(KeyError, s.get_key, wr) + self.assertNotIn(rd, s.get_map()) + self.assertNotIn(wr, s.get_map()) def test_fileno(self): s = self.SELECTOR()