Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from typing import TYPE_CHECKING, Optional

from threading import Thread
from .basic import Basic

from pieces._vendor.pieces_os_client.models.allocation_status_enum import AllocationStatusEnum
Expand Down Expand Up @@ -53,20 +53,27 @@ def _on_login_connect(self):
thread: The thread handling the login process.
timeout: The maximum time to wait for the login process.
"""
self.connect(True)
self.connect()

def login(self, connect_after_login=True, timeout=120):
"""
Logs the user into the OS and optionally connects to the cloud.

Args:
connect_after_login: A flag indicating if the user should connect to the cloud after login (default is True).
timeout: The maximum time to wait for the login process (default is 120 seconds).
connect_after_login: A flag indicating if the user should connect to the cloud after login (default is True).
timeout: The maximum time to wait for the login process (default is 120 seconds).
"""
thread = self.pieces_client.os_api.sign_into_os(async_req=True)
result = {}

def target():
result['user'] = self.pieces_client.os_api.sign_into_os()

thread = Thread(target=target)
thread.start()
thread.join(timeout)

if connect_after_login:
user = thread.get(timeout)
self.user_profile = user
self.user_profile = result.get('user')
self._on_login_connect()

def logout(self):
Expand All @@ -84,8 +91,20 @@ def connect(self, async_req = False):
"""
if not self.user_profile:
raise PermissionError("You must be logged in to use this feature")
self.on_user_callback(self.user_profile, True) # Set the connecting to cloud bool to true
self.pieces_client.allocations_api.allocations_connect_new_cloud(self.user_profile,async_req=async_req)
self.on_user_callback(
self.user_profile, True
) # Set the connecting to cloud bool to true
if async_req:
thread = Thread(
target=self.pieces_client.allocations_api.allocations_connect_new_cloud,
args=(self.user_profile,),
)
thread.start()
else:
self.pieces_client.allocations_api.allocations_connect_new_cloud(
self.user_profile
)


def disconnect(self):
"""
Expand Down
11 changes: 1 addition & 10 deletions src/pieces/_vendor/pieces_os_client/wrapper/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -188,7 +188,7 @@ def model_id(self):
def model_name(self, model):
models = self.get_models()
if model not in models:
raise ValueError("Not a vaild model name, the available models are"
raise ValueError("Not a valid model name, the available models are"
f"{', '.join(models.keys())}")
self._model_name = model
self._model_id = models[model]
Expand Down Expand Up @@ -282,15 +282,6 @@ def pieces_os_installer(self, callback: Callable[[DownloadModel], None]) -> PosI
"""
return PosInstaller(callback, self.app_name)

def pool(self, api_call, args):
"""
call the api async without stopping the main thread
Create thread pool on first request
avoids instantiating unused threadpool for blocking clients.
return the ThreadPool created
"""
return self.api_client.pool.apply_async(api_call, args)


# Register the function to be called on exit
atexit.register(PiecesClient.close)
9 changes: 8 additions & 1 deletion src/pieces/command_interface/auth_commands.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
import argparse
from pieces._vendor.pieces_os_client.models.allocation_status_enum import (
AllocationStatusEnum,
)
from pieces.base_command import BaseCommand
from pieces.urls import URLs
from pieces.settings import Settings
Expand Down Expand Up @@ -40,9 +43,13 @@ def execute(self, **kwargs) -> int:
Settings.pieces_client.user_api.user_snapshot().user
)
if Settings.pieces_client.user.user_profile:
status = Settings.pieces_client.user.cloud_status or AllocationStatusEnum.DISCONNECTED
Settings.logger.print(
f"Signed in as {Settings.pieces_client.user.name}\nemail: {Settings.pieces_client.user.email}"
f"Signed in as {Settings.pieces_client.user.name}\nEmail: {Settings.pieces_client.user.email}\nCloud status: {status.value.title()}"
)
if status == AllocationStatusEnum.DISCONNECTED:
Settings.logger.print("Connecting to the Pieces Cloud...")
Settings.pieces_client.user.connect()
return 0
try:
Settings.pieces_client.user.login(True)
Expand Down
4 changes: 3 additions & 1 deletion src/pieces/core/assets_command.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,9 +47,11 @@ def wrapper(*args, **kwargs):
try:
if AssetsCommands.current_asset is None:
raise ValueError("No material selected")
AssetsCommands.current_asset.asset # Check if the current asset is vaild
AssetsCommands.current_asset.asset # Check if the current asset is valid
except (ValueError, NotFoundException):
ListCommand.list_assets()
if AssetsCommands.current_asset is None:
return
return func(asset=AssetsCommands.current_asset, *args, **kwargs)

return wrapper
Expand Down
Loading