From 6f300b00795dd76ced62407f9eb065778a8bcbe6 Mon Sep 17 00:00:00 2001
From: Divyansh Singh <divyanshanandsingh@gmail.com>
Date: Fri, 24 Mar 2023 02:39:37 +0530
Subject: [PATCH] feat: constants support

---
 borgmatic/config/load.py              | 13 +++++++++++--
 borgmatic/config/schema.yaml          | 11 +++++++++++
 tests/integration/config/test_load.py | 17 ++++++++++++++++-
 3 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/borgmatic/config/load.py b/borgmatic/config/load.py
index 04461af..e379537 100644
--- a/borgmatic/config/load.py
+++ b/borgmatic/config/load.py
@@ -81,7 +81,8 @@ class Include_constructor(ruamel.yaml.SafeConstructor):
 def load_configuration(filename):
     '''
     Load the given configuration file and return its contents as a data structure of nested dicts
-    and lists.
+    and lists. Also, replace any "{constant}" strings with the value of the "constant" key in the
+    "constants" section of the configuration file.
 
     Raise ruamel.yaml.error.YAMLError if something goes wrong parsing the YAML, or RecursionError
     if there are too many recursive includes.
@@ -98,7 +99,15 @@ def load_configuration(filename):
     yaml = ruamel.yaml.YAML(typ='safe')
     yaml.Constructor = Include_constructor_with_include_directory
 
-    return yaml.load(open(filename))
+    with open(filename) as f:
+        file_contents = f.read()
+        config = yaml.load(file_contents)
+        if config and 'constants' in config:
+            for key, value in config['constants'].items():
+                file_contents = file_contents.replace(f'{{{key}}}', str(value))
+            config = yaml.load(file_contents)
+            del config['constants']
+        return config
 
 
 DELETED_NODE = object()
diff --git a/borgmatic/config/schema.yaml b/borgmatic/config/schema.yaml
index d4d57ab..c01f3db 100644
--- a/borgmatic/config/schema.yaml
+++ b/borgmatic/config/schema.yaml
@@ -3,6 +3,17 @@ required:
     - location
 additionalProperties: false
 properties:
+    constants:
+        type: object
+        description: |
+            Constants to use in the configuration file. All occurences of the
+            constant name within culy braces will be replaced with the value.
+            For example, if you have a constant named "hostname" with the value
+            "myhostname", then the string "{hostname}" will be replaced with
+            "myhostname" in the configuration file.
+        example:
+            hostname: myhostname
+            prefix: myprefix
     location:
         type: object
         description: |
diff --git a/tests/integration/config/test_load.py b/tests/integration/config/test_load.py
index e1ecc8a..5b8ee9b 100644
--- a/tests/integration/config/test_load.py
+++ b/tests/integration/config/test_load.py
@@ -10,8 +10,23 @@ from borgmatic.config import load as module
 
 def test_load_configuration_parses_contents():
     builtins = flexmock(sys.modules['builtins'])
-    builtins.should_receive('open').with_args('config.yaml').and_return('key: value')
+    config_file = io.StringIO('key: value')
+    config_file.name = 'config.yaml'
+    builtins.should_receive('open').with_args('config.yaml').and_return(config_file)
+    assert module.load_configuration('config.yaml') == {'key': 'value'}
 
+
+def test_load_configuration_replaces_constants():
+    builtins = flexmock(sys.modules['builtins'])
+    config_file = io.StringIO(
+        '''
+        constants:
+            key: value
+        key: {key}
+        '''
+    )
+    config_file.name = 'config.yaml'
+    builtins.should_receive('open').with_args('config.yaml').and_return(config_file)
     assert module.load_configuration('config.yaml') == {'key': 'value'}