Adding a param in checkout that allows init_branch (#522)

* Adding a param in checkout that allows init_branch

* Fix integration_client types for checkout

* Remove redundant call for list_branches

* Add docs + modify docstring

* Fix typo

* Using kwargs instead of args in code for checkout

docs/examples/api-reference/client/branch.py

@@ -55,4 +55,12 @@ def unused_checkout():
 
         # all subsequent operations will be on `mybranch`
         client.commit(...)
+
+        # create and change active branch from `mybranch` to `mybranch2`
+        # docsnip-highlight start
+        client.checkout("mybranch2", init=True)
+        assert client.branch() == "mybranch2"
+        # docsnip-highlight end
+
+        # all subsequent operations will be on `mybranch2`
         # /docsnip

docs/pages/api-reference/client/checkout.md

@@ -13,11 +13,11 @@ Sets the client to point to the given branch.
 <Expandable title="name" type="str">
 The name of the branch that the client should start pointing to. All subsequent
 operations (e.g. `commit`, `query`) will be directed to this branch.
-
-Note that `checkout` doesn't validate that the `name` points to a real branch.
-Instead, it just changes the local state of the client. If the branch doesn't 
-exist, subsequent branch operations will fail, not the `checkout` itself.
 </Expandable>
+<Expandable title="init" type="bool" defaultVal="False">
+If true, creates a new empty branch if the `name` is not found within the branches synced with Fennel
+</Expandable>
+
 
 <pre snippet="api-reference/client/branch#checkout" status="success"
     message="Changing client to point to 'mybranch'">
@@ -29,4 +29,9 @@ exist, subsequent branch operations will fail, not the `checkout` itself.
 
 :::info
 If not specified via explicit `checkout`, by default, clients point to the 'main' branch.
+:::
+
+:::info
+Note that `checkout` doesn't validate that the `name` points to a real branch by default. Instead, it just changes the local state of the client. If the branch doesn't 
+exist, subsequent branch operations will fail, not the `checkout` itself. However, when `init` is set to `True`, `checkout` will first create the branch if a real branch is not found and subsequently point to it.
 :::

fennel/client/client.py

@@ -430,14 +430,17 @@ def list_branches(self) -> List[str]:
             )
         return [branch.get("name") for branch in resp["branches"]]
 
-    def checkout(self, name: str):
+    def checkout(self, name: str, init: bool = False):
         """Checkouts the client to another branch.
 
         Parameters:
         ----------
         name (str): The name of the branch to checkout.
+        init (bool): If the branch should be created if it doesn't exist
 
         """
+        if init and name not in self.list_branches():
+            self.init_branch(name)
         self._branch = name
 
     def branch(self) -> str:

fennel/client_tests/branches/test_create.py

@@ -43,6 +43,22 @@ def test_simple_create(client):
     client.commit(message="msg")
 
 
+@pytest.mark.integration
+@mock
+def test_simple_create_with_checkout(client):
+    branch_list = client.list_branches()
+    assert len(branch_list) == 1
+
+    client.checkout("checkout-test")
+    branch_list = client.list_branches()
+    assert len(branch_list) == 1
+
+    client.checkout("checkout-test", init=True)
+    assert client.branch() == "checkout-test"
+    branch_list = client.list_branches()
+    assert len(branch_list) == 2
+
+
 @pytest.mark.integration
 @mock
 def test_duplicate_create(client):

fennel/testing/integration_client.py

@@ -86,7 +86,9 @@ def commit(
         time.sleep(10)
         return resp
 
-    def checkout(self, name: str):
+    def checkout(self, name: str, init: bool = False):
+        if init and name not in self.list_branches():
+            self.init_branch(name)
         self._branch = name
         fennel.datasets.datasets.dataset_lookup = partial(lookup_wrapper, name, False)  # type: ignore
 

fennel/testing/mock_client.py

@@ -438,7 +438,9 @@ def delete_branch(self, name: str):
     def list_branches(self) -> List[str]:
         return list(self.branches_map.keys())
 
-    def checkout(self, name: str):
+    def checkout(self, name: str, init: bool = False):
+        if init and name not in self.list_branches():
+            self.init_branch(name)
         self._branch = name
 
     # ----------------------- Secret API's -----------------------------------