Allow to call the root() BQL function with only one argument

Make the second argument default to 1.
beancount · Jun 10, 2024 · 5625b83 · 5625b83
1 parent e39b7d7
commit 5625b83
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/beanquery/query_env.py b/beanquery/query_env.py
@@ -249,8 +249,9 @@ def today():
 
 # Operations on accounts.
 
+@function([str], str)
 @function([str, int], str)
-def root(acc, n):
+def root(acc, n=1):
     """Get the root name(s) of the account."""
     return account.root(n, acc)
 

diff --git a/beanquery/query_env_test.py b/beanquery/query_env_test.py
@@ -263,6 +263,11 @@ def test_interval(self):
         self.assertResult('interval("1 baz")', None)
         self.assertResult('interval("A days")', None)
 
+    def test_root(self):
+        self.assertResult('root("Assets:Foo:Bar", 2)', 'Assets:Foo')
+        self.assertResult('root("Assets:Foo:Bar", 1)', 'Assets')
+        self.assertResult('root("Assets:Foo:Bar")', 'Assets')
+
     def test_func_meta(self):
         # use the loader to have the pad transaction inserted
         entries, _, options = loader.load_string('''