Preserve previous config data after update (#680)

* Add config updating mechanism

* Update tests

* Fix version not updating

streamrip/config.py

@@ -1,6 +1,7 @@
-"""A config class that manages arguments between the config file and CLI."""
+"""Classes and functions that manage config state."""
 
 import copy
+import functools
 import logging
 import os
 import shutil
@@ -19,6 +20,10 @@
 CURRENT_CONFIG_VERSION = "2.0.6"
 
 
+class OutdatedConfigError(Exception):
+    pass
+
+
 @dataclass(slots=True)
 class QobuzConfig:
     use_auth_token: bool
@@ -262,7 +267,7 @@ def from_toml(cls, toml_str: str):
         # TODO: handle the mistake where Windows people forget to escape backslash
         toml = parse(toml_str)
         if (v := toml["misc"]["version"]) != CURRENT_CONFIG_VERSION:  # type: ignore
-            raise Exception(
+            raise OutdatedConfigError(
                 f"Need to update config from {v} to {CURRENT_CONFIG_VERSION}",
             )
 
@@ -367,6 +372,26 @@ def save_file(self):
             self.file.update_toml()
             toml_file.write(dumps(self.file.toml))
 
+    @staticmethod
+    def _update_file(old_path: str, new_path: str):
+        """Updates the current config based on a newer config `new_toml`."""
+        with open(new_path) as new_conf:
+            new_toml = parse(new_conf.read())
+
+        toml_set_user_defaults(new_toml)
+
+        with open(old_path) as old_conf:
+            old_toml = parse(old_conf.read())
+
+        update_config(old_toml, new_toml)
+
+        with open(old_path, "w") as f:
+            f.write(dumps(new_toml))
+
+    @classmethod
+    def update_file(cls, path: str):
+        cls._update_file(path, BLANK_CONFIG_PATH)
+
     @classmethod
     def defaults(cls):
         return cls(BLANK_CONFIG_PATH)
@@ -384,9 +409,65 @@ def set_user_defaults(path: str, /):
 
     with open(path) as f:
         toml = parse(f.read())
+
+    toml_set_user_defaults(toml)
+
+    with open(path, "w") as f:
+        f.write(dumps(toml))
+
+
+def toml_set_user_defaults(toml: TOMLDocument):
     toml["downloads"]["folder"] = DEFAULT_DOWNLOADS_FOLDER  # type: ignore
     toml["database"]["downloads_path"] = DEFAULT_DOWNLOADS_DB_PATH  # type: ignore
     toml["database"]["failed_downloads_path"] = DEFAULT_FAILED_DOWNLOADS_DB_PATH  # type: ignore
     toml["youtube"]["video_downloads_folder"] = DEFAULT_YOUTUBE_VIDEO_DOWNLOADS_FOLDER  # type: ignore
-    with open(path, "w") as f:
-        f.write(dumps(toml))
+
+
+def _get_dict_keys_r(d: dict) -> set[tuple]:
+    """Get all possible key combinations in nested dicts.
+
+    See tests/test_config.py for example.
+    """
+    keys = d.keys()
+    ret = set()
+    for cur in keys:
+        val = d[cur]
+        if isinstance(val, dict):
+            ret.update((cur, *remaining) for remaining in _get_dict_keys_r(val))
+        else:
+            ret.add((cur,))
+    return ret
+
+
+def _nested_get(dictionary, *keys, default=None):
+    return functools.reduce(
+        lambda d, key: d.get(key, default) if isinstance(d, dict) else default,
+        keys,
+        dictionary,
+    )
+
+
+def _nested_set(dictionary, *keys, val):
+    """Nested set. Throws exception if keys are invalid."""
+    assert len(keys) > 0
+    final = functools.reduce(lambda d, key: d.get(key), keys[:-1], dictionary)
+    final[keys[-1]] = val
+
+
+def update_config(old_with_data: dict, new_without_data: dict):
+    """Used to update config when a new config version is detected.
+
+    All data associated with keys that are shared between the old and
+    new configs are copied from old to new. The remaining keep their default value.
+
+    Assumes that new_without_data contains default config values of the
+    latest version.
+    """
+    old_keys = _get_dict_keys_r(old_with_data)
+    new_keys = _get_dict_keys_r(new_without_data)
+    common = old_keys.intersection(new_keys)
+    common.discard(("misc", "version"))
+
+    for k in common:
+        old_val = _nested_get(old_with_data, *k)
+        _nested_set(new_without_data, *k, val=old_val)

streamrip/rip/cli.py

@@ -17,7 +17,7 @@
 from rich.traceback import install
 
 from .. import __version__, db
-from ..config import DEFAULT_CONFIG_PATH, Config, set_user_defaults
+from ..config import DEFAULT_CONFIG_PATH, Config, OutdatedConfigError, set_user_defaults
 from ..console import console
 from .main import Main
 
@@ -116,6 +116,11 @@ def rip(ctx, config_path, folder, no_db, quality, codec, no_progress, verbose):
 
     try:
         c = Config(config_path)
+    except OutdatedConfigError as e:
+        console.print(e)
+        console.print("Auto-updating config file...")
+        Config.update_file(config_path)
+        c = Config(config_path)
     except Exception as e:
         console.print(
             f"Error loading config from [bold cyan]{config_path}[/bold cyan]: {e}\n"

tests/test_config.py

@@ -1,10 +1,14 @@
+import os
 import shutil
 
 import pytest
+import tomlkit
 
 from streamrip.config import *
+from streamrip.config import _get_dict_keys_r, _nested_set
 
 SAMPLE_CONFIG = "tests/test_config.toml"
+OLD_CONFIG = "tests/test_config_old.toml"
 
 
 # Define a fixture to create a sample ConfigData instance for testing
@@ -26,6 +30,98 @@ def sample_config() -> Config:
     return config
 
 
+def test_get_keys_r():
+    d = {
+        "key1": {
+            "key2": {
+                "key3": 1,
+                "key4": 1,
+            },
+            "key6": [1, 2],
+            5: 1,
+        }
+    }
+    res = _get_dict_keys_r(d)
+    print(res)
+    assert res == {
+        ("key1", "key2", "key3"),
+        ("key1", "key2", "key4"),
+        ("key1", "key6"),
+        ("key1", 5),
+    }
+
+
+def test_safe_set():
+    d = {
+        "key1": {
+            "key2": {
+                "key3": 1,
+                "key4": 1,
+            },
+            "key6": [1, 2],
+            5: 1,
+        }
+    }
+    _nested_set(d, "key1", "key2", "key3", val=5)
+    assert d == {
+        "key1": {
+            "key2": {
+                "key3": 5,
+                "key4": 1,
+            },
+            "key6": [1, 2],
+            5: 1,
+        }
+    }
+
+
+def test_config_update():
+    old = {
+        "downloads": {"folder": "some_path", "use_service": True},
+        "qobuz": {"email": "asdf@gmail.com", "password": "test"},
+        "legacy_conf": {"something": 1, "other": 2},
+    }
+    new = {
+        "downloads": {"folder": "", "use_service": False, "keep_artwork": True},
+        "qobuz": {"email": "", "password": ""},
+        "tidal": {"email": "", "password": ""},
+    }
+    update_config(old, new)
+    assert new == {
+        "downloads": {"folder": "some_path", "use_service": True, "keep_artwork": True},
+        "qobuz": {"email": "asdf@gmail.com", "password": "test"},
+        "tidal": {"email": "", "password": ""},
+    }
+
+
+def test_config_throws_outdated():
+    with pytest.raises(Exception, match="update"):
+        _ = Config(OLD_CONFIG)
+
+
+def test_config_file_update():
+    tmp_conf = "tests/test_config_old2.toml"
+    shutil.copy("tests/test_config_old.toml", tmp_conf)
+    Config._update_file(tmp_conf, SAMPLE_CONFIG)
+
+    with open(tmp_conf) as f:
+        s = f.read()
+        toml = tomlkit.parse(s)  # type: ignore
+
+    assert toml["downloads"]["folder"] == "old_value"  # type: ignore
+    assert toml["downloads"]["source_subdirectories"] is True  # type: ignore
+    assert toml["downloads"]["concurrency"] is True  # type: ignore
+    assert toml["downloads"]["max_connections"] == 6  # type: ignore
+    assert toml["downloads"]["requests_per_minute"] == 60  # type: ignore
+    assert toml["cli"]["text_output"] is True  # type: ignore
+    assert toml["cli"]["progress_bars"] is True  # type: ignore
+    assert toml["cli"]["max_search_results"] == 100  # type: ignore
+    assert toml["misc"]["version"] == "2.0.6"  # type: ignore
+    assert "YouTubeVideos" in str(toml["youtube"]["video_downloads_folder"])
+    # type: ignore
+    os.remove("tests/test_config_old2.toml")
+
+
 def test_sample_config_data_properties(sample_config_data):
     # Test the properties of ConfigData
     assert sample_config_data.modified is False  # Ensure initial state is not modified