overload GetPlays for Infoset and Action

Author: d-kad

src/games/game.h

@@ -476,6 +476,10 @@ class GameRep : public BaseGameRep {
 
   /// Returns the set of terminal nodes which are descendants of node
   virtual std::vector<GameNode> GetPlays(GameNode node) const = 0;
+  /// Returns the set of terminal nodes which are descendants of members of an infoset
+  virtual std::vector<GameNode> GetPlays(GameInfoset infoset) const = 0;
+  /// Returns the set of terminal nodes which are descendants of members of an action
+  virtual std::vector<GameNode> GetPlays(GameAction action) const = 0;
 
   /// Returns true if the game is perfect recall.  If not,
   /// a pair of violating information sets is returned in the parameters.

src/games/gameagg.h

@@ -105,6 +105,13 @@ class GameAGGRep : public GameRep {
   size_t NumNonterminalNodes() const override { throw UndefinedException(); }
   /// Returns the set of terminal nodes which are descendants of node
   std::vector<GameNode> GetPlays(GameNode node) const override { throw UndefinedException(); }
+  /// Returns the set of terminal nodes which are descendants of members of an infoset
+  std::vector<GameNode> GetPlays(GameInfoset infoset) const override
+  {
+    throw UndefinedException();
+  }
+  /// Returns the set of terminal nodes which are descendants of members of an action
+  std::vector<GameNode> GetPlays(GameAction action) const override { throw UndefinedException(); }
   //@}
 
   /// @name General data access

src/games/gamebagg.h

@@ -113,6 +113,13 @@ class GameBAGGRep : public GameRep {
   size_t NumNonterminalNodes() const override { throw UndefinedException(); }
   /// Returns the set of terminal nodes which are descendants of node
   std::vector<GameNode> GetPlays(GameNode node) const override { throw UndefinedException(); }
+  /// Returns the set of terminal nodes which are descendants of members of an infoset
+  std::vector<GameNode> GetPlays(GameInfoset infoset) const override
+  {
+    throw UndefinedException();
+  }
+  /// Returns the set of terminal nodes which are descendants of members of an action
+  std::vector<GameNode> GetPlays(GameAction action) const override { throw UndefinedException(); }
   //@}
 
   /// @name General data access

src/games/gametable.h

@@ -94,6 +94,13 @@ class GameTableRep : public GameExplicitRep {
   size_t NumNonterminalNodes() const override { throw UndefinedException(); }
   /// Returns the set of terminal nodes which are descendants of node
   std::vector<GameNode> GetPlays(GameNode node) const override { throw UndefinedException(); }
+  /// Returns the set of terminal nodes which are descendants of members of an infoset
+  std::vector<GameNode> GetPlays(GameInfoset infoset) const override
+  {
+    throw UndefinedException();
+  }
+  /// Returns the set of terminal nodes which are descendants of members of an action
+  std::vector<GameNode> GetPlays(GameAction action) const override { throw UndefinedException(); }
   //@}
 
   /// @name Outcomes

src/games/gametree.cc

@@ -1080,6 +1080,31 @@ std::vector<GameNode> GameTreeRep::GetPlays(GameNode node) const
   return consistent_plays_copy;
 }
 
+std::vector<GameNode> GameTreeRep::GetPlays(GameInfoset infoset) const
+{
+  std::vector<GameNode> plays;
+  auto members = infoset->GetMembers();
+
+  for (const GameNode &node : members) {
+    std::vector<GameNode> member_plays = GetPlays(node);
+    plays.insert(plays.end(), member_plays.begin(), member_plays.end());
+  }
+  return plays;
+}
+
+std::vector<GameNode> GameTreeRep::GetPlays(GameAction action) const
+{
+  std::vector<GameNode> plays;
+  auto infoset = action->GetInfoset();
+  auto members = infoset->GetMembers();
+
+  for (const GameNode &node : members) {
+    std::vector<GameNode> child_plays = GetPlays(node->GetChild(action));
+    plays.insert(plays.end(), child_plays.begin(), child_plays.end());
+  }
+  return plays;
+}
+
 void GameTreeRep::DeleteOutcome(const GameOutcome &p_outcome)
 {
   IncrementVersion();

src/games/gametree.h

@@ -146,6 +146,8 @@ class GameTreeRep : public GameExplicitRep {
   void SetOutcome(GameNode, const GameOutcome &p_outcome) override;
 
   std::vector<GameNode> GetPlays(GameNode node) const override;
+  std::vector<GameNode> GetPlays(GameInfoset infoset) const override;
+  std::vector<GameNode> GetPlays(GameAction action) const override;
 
   Game CopySubgame(GameNode) const override;
   //@}

src/pygambit/gambit.pxd

@@ -209,6 +209,8 @@ cdef extern from "games/game.h":
         c_Rational GetMaxPayoff() except +
         c_Rational GetMaxPayoff(c_GamePlayer) except +
         stdvector[c_GameNode] GetPlays(c_GameNode) except +
+        stdvector[c_GameNode] GetPlays(c_GameInfoset) except +
+        stdvector[c_GameNode] GetPlays(c_GameAction) except +
         bool IsPerfectRecall() except +
 
         c_GameInfoset AppendMove(c_GameNode, c_GamePlayer, int) except +ValueError

src/pygambit/game.pxi

@@ -22,7 +22,6 @@
 import io
 import itertools
 import pathlib
-import warnings
 
 import numpy as np
 import scipy.stats
@@ -2014,11 +2013,20 @@ class Game:
             raise UndefinedOperationError("Cannot delete the only strategy for a player")
         self.game.deref().DeleteStrategy(resolved_strategy.strategy)
 
-    def get_plays(self, node: typing.Union[Node, str]) -> typing.List[Node]:
-        resolved_node = cython.cast(Node, self._resolve_node(node, "get_plays", "node_obj"))
+    def get_plays(self, obj: typing.Union[Node, Infoset, Action]) -> typing.List[Node]:
 
-        plays = []
-        for item in self.game.deref().GetPlays(resolved_node.node):
-            plays.append(Node.wrap(item))
+        c_plays = cython.declare(stdvector[c_GameNode])
 
-        return plays
+        if isinstance(obj, Node):
+            obj = cython.cast(Node, obj)
+            c_plays = self.game.deref().GetPlays(obj.node)
+        elif isinstance(obj, Infoset):
+            obj = cython.cast(Infoset, obj)
+            c_plays = self.game.deref().GetPlays(obj.infoset)
+        elif isinstance(obj, Action):
+            obj = cython.cast(Action, obj)
+            c_plays = self.game.deref().GetPlays(obj.action)
+        else:
+            raise TypeError("The object needs to be either Node, Infoset, or Action")
+
+        return [Node.wrap(item) for item in c_plays]