reuse type

Author: nazarfil

openhexa/sdk/workspaces/current_workspace.py

@@ -5,6 +5,7 @@
 
 import os
 from dataclasses import fields, make_dataclass
+from typing import List
 from warnings import warn
 
 from openhexa.utils import stringcase
@@ -21,21 +22,7 @@
     S3Connection,
 )
 from openhexa.cli.api import OpenHexaClient
-
-
-class Country:
-    """Represents a country with its code, name, alpha3 code and flag."""
-
-    def __init__(self, code: str, name: str, alpha3: str, flag: str):
-        """Initialize a Country object."""
-        self.code = code
-        self.name = name
-        self.alpha3 = alpha3
-        self.flag = flag
-
-    def __repr__(self):
-        """Return a string representation of the Country object."""
-        return f"Country(code={self.code}, name={self.name}, alpha3={self.alpha3}, flag={self.flag})"
+from ...cli.graphql.graphql_client import GetCountriesWorkspaceCountries
 
 
 class WorkspaceConfigError(Exception):
@@ -76,7 +63,7 @@ def slug(self) -> str:
             raise WorkspaceConfigError("The workspace slug is not available in this environment.")
 
     @property
-    def countries(self) -> list[Country]:
+    def countries(self) -> List[GetCountriesWorkspaceCountries]:
         """The countries of the workspace."""
         try:
             return OpenHexaClient().get_countries(workspace_slug=self.slug).workspace.countries