fixes

katrain/core/game_node.py

@@ -2,13 +2,19 @@ import copy
 import random
 from typing import Dict, List, Optional, Tuple
 
-from katrain.core.constants import HOMEPAGE, VERSION, SGF_COMMENTS_MARKER, SGF_SEPARATOR_MARKER, SGF_NOTE_MARKER, PROGRAM_NAME
+from katrain.core.constants import (
+    VERSION,
+    SGF_INTERNAL_COMMENTS_MARKER,
+    SGF_SEPARATOR_MARKER,
+    PROGRAM_NAME,
+)
 from katrain.core.lang import i18n
 from katrain.core.sgf_parser import Move, SGFNode
 from katrain.core.utils import evaluation_class, var_to_grid
 from katrain.gui.style import INFO_PV_COLOR
 import base64
 import gzip
+import json
 
 
 class GameNode(SGFNode):
@@ -16,8 +22,6 @@ class GameNode(SGFNode):
 
     def __init__(self, parent=None, properties=None, move=None):
         super().__init__(parent=parent, properties=properties, move=move)
-        self.ownership = None
-        self.policy = None
         self.auto_undo = None  # None = not analyzed. False: not undone (good move). True: undone (bad move)
         self.ai_thoughts = ""
         self.note = ""
@@ -31,45 +35,41 @@ class GameNode(SGFNode):
     def add_list_property(self, property: str, values: List):
         if property == "KT":
             try:
-                import bson
-
-                self.analysis = bson.loads(gzip.decompress(base64.standard_b64decode(values[0])))
+                self.analysis = json.loads(gzip.decompress(base64.standard_b64decode(values[0])))
                 self.analysis_loaded = True
             except Exception as e:
-                print(e)
+                print(f"Error in loading analysis: {e}")
         elif property == "C":
             comments = [  # strip out all previously auto generated comments
-                c for v in values for c in v.split(SGF_SEPARATOR_MARKER)
+                c
+                for v in values
+                for c in v.split(SGF_SEPARATOR_MARKER)
+                if c.strip() and SGF_INTERNAL_COMMENTS_MARKER not in c
             ]
-            self.note = "".join(c for c in comments if SGF_NOTE_MARKER in c)
-            super().add_list_property(
-                property,
-                [c for c in comments if c.strip() and SGF_COMMENTS_MARKER not in c and SGF_NOTE_MARKER not in c],
-            )
+            self.note = "".join(comments)  # no super call intended, just save as note to be editable
         else:
             return super().add_list_property(property, values)
 
     def clear_analysis(self):
         self.analysis_visits_requested = 0
-        self.analysis = {"moves": {}, "root": None, "completed": False}
+        self.analysis = {"moves": {}, "root": None, "ownership": None, "policy": None, "completed": False}
 
     def sgf_properties(
         self, save_comments_player=None, save_comments_class=None, eval_thresholds=None, save_analysis=False
     ):
         properties = copy.copy(super().sgf_properties())
         note = self.note.strip()
-        print(save_analysis, self.analysis_complete)
         if save_analysis and self.analysis_complete:
             try:
-                import bson
-
                 analysis = copy.deepcopy(self.analysis)
                 for movedict in analysis["moves"].values():
                     if "ownership" in movedict:
                         del movedict["ownership"]
-                properties["KT"] = [base64.standard_b64encode(gzip.compress(bson.dumps(analysis))).decode("utf-8")]
+                properties["KT"] = [
+                    base64.standard_b64encode(gzip.compress(json.dumps(analysis).encode("utf-8"))).decode("utf-8")
+                ]
             except Exception as e:
-                print(e)
+                print(f"Error in saving analysis: {e}")
         if self.points_lost and save_comments_class is not None and eval_thresholds is not None:
             show_class = save_comments_class[evaluation_class(self.points_lost, eval_thresholds)]
         else:
@@ -93,12 +93,12 @@ class GameNode(SGFNode):
             comments.append(self.comment(sgf=True, interactive=False))
         if self.is_root:
             comments = [
-                i18n._("SGF start message") + SGF_COMMENTS_MARKER,
+                i18n._("SGF start message") + SGF_INTERNAL_COMMENTS_MARKER + "\n",
                 *comments,
-                f"SGF with generated by {PROGRAM_NAME} {VERSION}{SGF_COMMENTS_MARKER}.",
+                f"\nSGF with generated by {PROGRAM_NAME} {VERSION}{SGF_INTERNAL_COMMENTS_MARKER}.\n",
             ]
         if note:
-            comments.append(f"\nNote: {self.note}{SGF_NOTE_MARKER}")
+            comments.append(f"{self.note}")
         if comments:
             properties["C"] = [SGF_SEPARATOR_MARKER.join(comments)]
         elif "C" in properties:
@@ -169,8 +169,8 @@ class GameNode(SGFNode):
                     move_dict["order"] = 999  # old moves to end
             for move_analysis in analysis_json["moveInfos"]:
                 self.update_move_analysis(move_analysis, move_analysis["move"])
-            self.ownership = analysis_json.get("ownership")
-            self.policy = analysis_json.get("policy")
+            self.analysis["ownership"] = analysis_json.get("ownership")
+            self.analysis["policy"] = analysis_json.get("policy")
             if not additional_moves:
                 self.analysis["root"] = analysis_json["rootInfo"]
             if self.parent and self.move:
@@ -183,6 +183,14 @@ class GameNode(SGFNode):
             is_normal_query = refine_move is None and not additional_moves
             self.analysis["completed"] = self.analysis["completed"] or (is_normal_query and not partial_result)
 
+    @property
+    def ownership(self):
+        return self.analysis.get("ownership")
+
+    @property
+    def policy(self):
+        return self.analysis.get("policy")
+
     @property
     def analysis_exists(self):
         return self.analysis["root"] is not None