From 3829d262c325b61b0004a314347f985008a9eb84 Mon Sep 17 00:00:00 2001
From: Niklas Fiekas <niklas.fiekas@backscattering.de>
Date: Mon, 17 Jun 2024 21:18:43 +0200
Subject: [PATCH] Factor out Tablebase.add_file()

---
 chess/syzygy.py | 36 ++++++++++++++++--------------------
 1 file changed, 16 insertions(+), 20 deletions(-)

diff --git a/chess/syzygy.py b/chess/syzygy.py
index e1fe07eb..0c7d4150 100644
--- a/chess/syzygy.py
+++ b/chess/syzygy.py
@@ -1511,27 +1511,23 @@ def add_directory(self, directory: str, *, load_wdl: bool = True, load_dtz: bool
 
         Returns the number of table files that were found.
         """
-        num = 0
         directory = os.path.abspath(directory)
-
-        for filename in os.listdir(directory):
-            path = os.path.join(directory, filename)
-            tablename, ext = os.path.splitext(filename)
-
-            if is_tablename(tablename, one_king=self.variant.one_king) and os.path.isfile(path):
-                if load_wdl:
-                    if ext == self.variant.tbw_suffix:
-                        num += self._open_table(self.wdl, WdlTable, path)
-                    elif "P" not in tablename and ext == self.variant.pawnless_tbw_suffix:
-                        num += self._open_table(self.wdl, WdlTable, path)
-
-                if load_dtz:
-                    if ext == self.variant.tbz_suffix:
-                        num += self._open_table(self.dtz, DtzTable, path)
-                    elif "P" not in tablename and ext == self.variant.pawnless_tbz_suffix:
-                        num += self._open_table(self.dtz, DtzTable, path)
-
-        return num
+        return sum(self.add_file(os.path.join(directory, filename), load_wdl=load_wdl, load_dtz=load_dtz) for filename in os.listdir(directory))
+
+    def add_file(self, path: str, *, load_wdl: bool = True, load_dtz: bool = True) -> int:
+        tablename, ext = os.path.splitext(os.path.basename(path))
+        if is_tablename(tablename, one_king=self.variant.one_king) and os.path.isfile(path):
+            if load_wdl:
+                if ext == self.variant.tbw_suffix:
+                    return self._open_table(self.wdl, WdlTable, path)
+                elif "P" not in tablename and ext == self.variant.pawnless_tbw_suffix:
+                    return self._open_table(self.wdl, WdlTable, path)
+            if load_dtz:
+                if ext == self.variant.tbz_suffix:
+                    return self._open_table(self.dtz, DtzTable, path)
+                elif "P" not in tablename and ext == self.variant.pawnless_tbz_suffix:
+                    return self._open_table(self.dtz, DtzTable, path)
+        return 0
 
     def probe_wdl_table(self, board: chess.Board) -> int:
         # Test for variant end.