import uuid from drf_spectacular.utils import extend_schema from rest_framework.viewsets import GenericViewSet # type: ignore from rest_framework.decorators import action # type: ignore from rest_framework import status # type: ignore from rest_framework.request import HttpRequest # type: ignore from rest_framework.response import Response # type: ignore from rest_framework.permissions import ( # type: ignore IsAdminUser, ) from django_core.mixins import BaseViewSetMixin from rest_framework.generics import get_object_or_404 # type: ignore from django.contrib.auth import get_user_model from django.db import transaction from core.apps.companies.serializers import ( CreateCompanySerializer, RetrieveCompanySerializer ) from core.apps.companies.models import ( CompanyModel, CompanyAccountModel, ) UserModel = get_user_model() class UserCompaniesView(BaseViewSetMixin, GenericViewSet): permission_classes = [IsAdminUser] action_permission_classes = {} action_permission_classes = { "list_company": RetrieveCompanySerializer, "create_company": CreateCompanySerializer, } @extend_schema( summary="Get list of companies", description="Get list of companies", ) @action(url_path="companies", detail=True, methods=["GET"]) def list_company( self, request: HttpRequest, pk: uuid.UUID, *args: object, **kwargs: object, ) -> Response: companies = CompanyModel.objects.filter(accounts__user__pk=pk) return Response( data=RetrieveCompanySerializer(instance=companies, many=True), status=status.HTTP_200_OK ) @extend_schema( summary="Create Company", description="Create Company", ) @action(url_path="companies", detail=True, methods=["POST"]) def create_company( self, request: HttpRequest, pk: uuid.UUID, *args: object, **kwargs: object, ) -> Response: with transaction.atomic(): ser = CreateCompanySerializer(data=request.data) # type: ignore ser.is_valid(raise_exception=True) company = ser.save() # type: ignore user = get_object_or_404(UserModel, pk=pk) account = CompanyAccountModel(company=company, user=user) account.save() return Response(data=ser.data, status=status.HTTP_201_CREATED)