From f1f450cedc6ed9d79bf45f1045fc933c403bd171 Mon Sep 17 00:00:00 2001
From: andygoblins <andygoblins@gmx.com>
Date: Sun, 9 Feb 2020 14:53:17 -0600
Subject: [PATCH] make Session a context manager

---
 .../python/example_scripts/simple_test.py     | 120 +++++++++---------
 bindings/python/gnucash_core.py               |   9 ++
 2 files changed, 67 insertions(+), 62 deletions(-)

diff --git a/bindings/python/example_scripts/simple_test.py b/bindings/python/example_scripts/simple_test.py
index 9bc21789e6..b8a1bed02e 100644
--- a/bindings/python/example_scripts/simple_test.py
+++ b/bindings/python/example_scripts/simple_test.py
@@ -7,81 +7,77 @@ from gnucash import Session, Account, Transaction, Split, GncNumeric
 
 FILE_1 = "/tmp/example.gnucash"
 
-session = Session("xml://%s" % FILE_1, is_new=True)
+with Session("xml://%s" % FILE_1, is_new=True) as session:
 
-book = session.book
-root_acct = Account(book)
-expenses_acct = Account(book)
-savings_acct = Account(book)
-opening_acct = Account(book)
-trans1 = Transaction(book)
-trans1.BeginEdit()
-trans2 = Transaction(book)
-trans2.BeginEdit()
+    book = session.book
+    root_acct = Account(book)
+    expenses_acct = Account(book)
+    savings_acct = Account(book)
+    opening_acct = Account(book)
+    trans1 = Transaction(book)
+    trans1.BeginEdit()
+    trans2 = Transaction(book)
+    trans2.BeginEdit()
 
-split1 = Split(book)
-split3 = Split(book)
-comm_table = book.get_table()
-cad = comm_table.lookup("CURRENCY", "CAD")
+    split1 = Split(book)
+    split3 = Split(book)
+    comm_table = book.get_table()
+    cad = comm_table.lookup("CURRENCY", "CAD")
 
-num1 = GncNumeric(4, 1)
-num2 = GncNumeric(100, 1)
+    num1 = GncNumeric(4, 1)
+    num2 = GncNumeric(100, 1)
 
-#Set new root account
-book.set_root_account(root_acct)
+    #Set new root account
+    book.set_root_account(root_acct)
 
-#Set up root account and add sub-accounts
-root_acct.SetName("Root")
-root_acct.SetType(13) #ACCT_TYPE_ROOT = 13
-root_acct.append_child(expenses_acct)
-root_acct.append_child(savings_acct)
-root_acct.append_child(opening_acct)
+    #Set up root account and add sub-accounts
+    root_acct.SetName("Root")
+    root_acct.SetType(13) #ACCT_TYPE_ROOT = 13
+    root_acct.append_child(expenses_acct)
+    root_acct.append_child(savings_acct)
+    root_acct.append_child(opening_acct)
 
-#Set up Expenses account
-expenses_acct.SetCommodity(cad)
-expenses_acct.SetName("Expenses")
-expenses_acct.SetType(9) #ACCT_TYPE_EXPENSE = 9
+    #Set up Expenses account
+    expenses_acct.SetCommodity(cad)
+    expenses_acct.SetName("Expenses")
+    expenses_acct.SetType(9) #ACCT_TYPE_EXPENSE = 9
 
-#Set up Savings account
-savings_acct.SetCommodity(cad)
-savings_acct.SetName("Savings")
-savings_acct.SetType(0) #ACCT_TYPE_BANK = 0
+    #Set up Savings account
+    savings_acct.SetCommodity(cad)
+    savings_acct.SetName("Savings")
+    savings_acct.SetType(0) #ACCT_TYPE_BANK = 0
 
-#Set up Opening Balance account
-opening_acct.SetCommodity(cad)
-opening_acct.SetName("Opening Balance")
-opening_acct.SetType(10) #ACCT_TYPE_EQUITY = 10
+    #Set up Opening Balance account
+    opening_acct.SetCommodity(cad)
+    opening_acct.SetName("Opening Balance")
+    opening_acct.SetType(10) #ACCT_TYPE_EQUITY = 10
 
-split1.SetValue(num1)
-split1.SetAccount(expenses_acct)
-split1.SetParent(trans1)
+    split1.SetValue(num1)
+    split1.SetAccount(expenses_acct)
+    split1.SetParent(trans1)
 
-split3.SetValue(num2)
-split3.SetAccount(savings_acct)
-split3.SetParent(trans2)
+    split3.SetValue(num2)
+    split3.SetAccount(savings_acct)
+    split3.SetParent(trans2)
 
-trans1.SetCurrency(cad)
-trans1.SetDate(14, 3, 2006)
-trans1.SetDescription("Groceries")
+    trans1.SetCurrency(cad)
+    trans1.SetDate(14, 3, 2006)
+    trans1.SetDescription("Groceries")
 
-trans2.SetCurrency(cad)
-trans2.SetDate(7, 11, 1995)
-trans2.SetDescription("Opening Savings Balance")
+    trans2.SetCurrency(cad)
+    trans2.SetDate(7, 11, 1995)
+    trans2.SetDescription("Opening Savings Balance")
 
-split2 = Split(book)
-split2.SetAccount(savings_acct)
-split2.SetParent(trans1)
-split2.SetValue(num1.neg())
+    split2 = Split(book)
+    split2.SetAccount(savings_acct)
+    split2.SetParent(trans1)
+    split2.SetValue(num1.neg())
 
-split4 = Split(book)
-split4.SetAccount(opening_acct)
-split4.SetParent(trans2)
-split4.SetValue(num2.neg())
+    split4 = Split(book)
+    split4.SetAccount(opening_acct)
+    split4.SetParent(trans2)
+    split4.SetValue(num2.neg())
 
 
-trans1.CommitEdit()
-trans2.CommitEdit()
-
-session.save()
-session.end()
-session.destroy()
+    trans1.CommitEdit()
+    trans2.CommitEdit()
\ No newline at end of file
diff --git a/bindings/python/gnucash_core.py b/bindings/python/gnucash_core.py
index 8cab524177..915fe06383 100644
--- a/bindings/python/gnucash_core.py
+++ b/bindings/python/gnucash_core.py
@@ -117,6 +117,15 @@ class Session(GnuCashCoreClass):
                 self.destroy()
                 raise
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        # Roll back changes on exception by not calling save. Only works for XMl backend.
+        if not exc_type:
+            self.save()
+        self.end()
+
     def raise_backend_errors(self, called_function="qof_session function"):
         """Raises a GnuCashBackendException if there are outstanding
         QOF_BACKEND errors.