From c1cfdffc9666f011a42155928891ccc2171abfe0 Mon Sep 17 00:00:00 2001
From: Dave Pearson <davep@davep.org>
Date: Thu, 3 Nov 2022 13:44:51 +0000
Subject: [PATCH] Change query_one so that it raises an error on multiple hits

As part of this, add a only_one method to DOMQuery.

Addresses #1096.
---
 CHANGELOG.md             |  4 ++++
 src/textual/css/query.py | 47 ++++++++++++++++++++++++++++++++++++++++
 src/textual/dom.py       |  5 +----
 tests/test_query.py      |  6 ++++-
 4 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e952920e36..71350ccefd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,11 +12,15 @@ and this project adheres to [Semantic Versioning](http://semver.org/).
 - Dropped support for mounting "named" and "anonymous" widgets via
   `App.mount` and `Widget.mount`. Both methods now simply take one or more
   widgets as positional arguments.
+- `DOMNode.query_one` now raises a `TooManyMatches` exception if there is
+  more than one matching node.
+  https://github.com/Textualize/textual/issues/1096
 
 ### Added
 
 - Added `init` param to reactive.watch
 - `CSS_PATH` can now be a list of CSS files https://github.com/Textualize/textual/pull/1079
+- Added `DOMQuery.only_one` https://github.com/Textualize/textual/issues/1096
 
 ## [0.3.0] - 2022-10-31
 
diff --git a/src/textual/css/query.py b/src/textual/css/query.py
index e0d5882046..30c060d416 100644
--- a/src/textual/css/query.py
+++ b/src/textual/css/query.py
@@ -42,6 +42,10 @@ class NoMatches(QueryError):
     """No nodes matched the query."""
 
 
+class TooManyMatches(QueryError):
+    """Too many nodes matched the query."""
+
+
 class WrongType(QueryError):
     """Query result was not of the correct type."""
 
@@ -208,6 +212,49 @@ def first(
         else:
             raise NoMatches(f"No nodes match {self!r}")
 
+    @overload
+    def only_one(self) -> Widget:
+        ...
+
+    @overload
+    def only_one(self, expect_type: type[ExpectType]) -> ExpectType:
+        ...
+
+    def only_one(
+        self, expect_type: type[ExpectType] | None = None
+    ) -> Widget | ExpectType:
+        """Get the *only* matching node.
+
+        Args:
+            expect_type (type[ExpectType] | None, optional): Require matched node is of this type,
+                or None for any type. Defaults to None.
+
+        Raises:
+            WrongType: If the wrong type was found.
+            TooManyMatches: If there is more than one matching node in the query.
+
+        Returns:
+            Widget | ExpectType: The matching Widget.
+        """
+        # Call on first to get the first item. Here we'll use all of the
+        # testing and checking it provides.
+        the_one = self.first(expect_type) if expect_type is not None else self.first()
+        try:
+            # Now see if we can access a subsequent item in the nodes. There
+            # should *not* be anything there, so we *should* get an
+            # IndexError. We *could* have just checked the length of the
+            # query, but the idea here is to do the check as cheaply as
+            # possible.
+            _ = self.nodes[1]
+            raise TooManyMatches(
+                "Call to only_one resulted in more than one matched node"
+            )
+        except IndexError:
+            # The IndexError was got, that's a good thing in this case. So
+            # we return what we found.
+            pass
+        return the_one
+
     @overload
     def last(self) -> Widget:
         ...
diff --git a/src/textual/dom.py b/src/textual/dom.py
index 2efbb14b0e..3521eab8d8 100644
--- a/src/textual/dom.py
+++ b/src/textual/dom.py
@@ -783,10 +783,7 @@ def query_one(
             query_selector = selector.__name__
         query: DOMQuery[Widget] = DOMQuery(self, filter=query_selector)
 
-        if expect_type is None:
-            return query.first()
-        else:
-            return query.first(expect_type)
+        return query.only_one() if expect_type is None else query.only_one(expect_type)
 
     def set_styles(self, css: str | None = None, **update_styles) -> None:
         """Set custom styles on this object."""
diff --git a/tests/test_query.py b/tests/test_query.py
index c48587ea19..8cc4de653b 100644
--- a/tests/test_query.py
+++ b/tests/test_query.py
@@ -1,7 +1,7 @@
 import pytest
 
 from textual.widget import Widget
-from textual.css.query import InvalidQueryFormat, WrongType, NoMatches
+from textual.css.query import InvalidQueryFormat, WrongType, NoMatches, TooManyMatches
 
 
 def test_query():
@@ -82,6 +82,10 @@ class App(Widget):
             helpbar,
         ]
         assert list(app.query("Widget.float").results(View)) == []
+        assert app.query_one("#widget1") == widget1
+        assert app.query_one("#widget1", Widget) == widget1
+        with pytest.raises(TooManyMatches):
+            _ = app.query_one(Widget)
 
         assert app.query("Widget.float")[0] == sidebar
         assert app.query("Widget.float")[0:2] == [sidebar, tooltip]