Skip to content

Commit

Permalink
Replace usage of Enum plus str mixin with StrEnum
Browse files Browse the repository at this point in the history
  • Loading branch information
lqc committed Apr 25, 2024
1 parent 6893f28 commit 7dfa7f7
Show file tree
Hide file tree
Showing 80 changed files with 221 additions and 224 deletions.
2 changes: 1 addition & 1 deletion docs/contributing/charts/how-to-add.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ sidebarTitle: "Charts"
Add a new type in [`mage_ai/data_preparation/models/widget/constants.py`](https://github.com/mage-ai/mage-ai/blob/master/mage_ai/data_preparation/models/widget/constants.py):

```python
class ChartType(str, Enum):
class ChartType(StrEnum):
# ...
PIE_CHART = 'pie_chart'
```
Expand Down
4 changes: 2 additions & 2 deletions mage_ai/ai/constants.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from enum import Enum
from enum import StrEnum


class LLMUseCase(str, Enum):
class LLMUseCase(StrEnum):
GENERATE_DOC_FOR_BLOCK = 'generate_doc_for_block'
GENERATE_DOC_FOR_PIPELINE = 'generate_doc_for_pipeline'
GENERATE_BLOCK_WITH_DESCRIPTION = 'generate_block_with_description'
Expand Down
8 changes: 4 additions & 4 deletions mage_ai/api/constants.py
Original file line number Diff line number Diff line change
@@ -1,20 +1,20 @@
from enum import Enum
from enum import StrEnum

from mage_ai.api.operations.constants import OperationType
from mage_ai.orchestration.db.models.oauth import Permission


class AttributeOperationType(str, Enum):
class AttributeOperationType(StrEnum):
QUERY = 'query'
READ = 'read'
WRITE = 'write'


class AttributeType(str, Enum):
class AttributeType(StrEnum):
ALL = '__*MAGE*__'


class AuthorizeStatusType(str, Enum):
class AuthorizeStatusType(StrEnum):
ALL = 'all'
FAILED = 'failed'
SUCCEEDED = 'succeeded'
Expand Down
4 changes: 2 additions & 2 deletions mage_ai/api/oauth_scope.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from enum import Enum
from enum import StrEnum


class OauthScope():
Expand All @@ -18,7 +18,7 @@ class OauthScope():
TOKEN_SCOPES = []


class OauthScopeType(str, Enum):
class OauthScopeType(StrEnum):
CLIENT_ALL = 'all'
CLIENT_INTERNAL = 'internal'
CLIENT_PRIVATE = 'private'
Expand Down
4 changes: 2 additions & 2 deletions mage_ai/api/operations/constants.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from enum import Enum
from enum import StrEnum

ALL = 'all'
CREATE = 'create'
Expand All @@ -20,7 +20,7 @@
COOKIE_PREFIX = '__COOKIE__'


class OperationType(str, Enum):
class OperationType(StrEnum):
ALL = ALL
CREATE = CREATE
DELETE = DELETE
Expand Down
2 changes: 1 addition & 1 deletion mage_ai/api/resources/PipelineResource.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import asyncio
import urllib.parse
from datetime import datetime, timedelta
from enum import Enum
from enum import StrEnum
from typing import Dict, List

from sqlalchemy import or_
Expand Down
4 changes: 2 additions & 2 deletions mage_ai/authentication/oauth/constants.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from enum import Enum
from enum import StrEnum
from typing import Optional

from mage_ai.settings import get_settings_value
Expand All @@ -10,7 +10,7 @@
GITHUB_STATE = '1337'


class ProviderName(str, Enum):
class ProviderName(StrEnum):
ACTIVE_DIRECTORY = 'active_directory'
AZURE_DEVOPS = 'azure_devops'
BITBUCKET = 'bitbucket'
Expand Down
4 changes: 2 additions & 2 deletions mage_ai/authentication/operation_history/constants.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
from enum import Enum
from enum import StrEnum

MAGE_OPERATION_HISTORY_DIRECTORY_DEFAULT = '.operation_history'
MAGE_OPERATION_HISTORY_DIRECTORY_ENVIRONMENT_VARIABLE_NAME = 'MAGE_OPERATION_HISTORY_DIRECTORY'


class ResourceType(str, Enum):
class ResourceType(StrEnum):
PIPELINE = 'pipeline'
10 changes: 5 additions & 5 deletions mage_ai/authentication/permissions/constants.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
from enum import Enum
from enum import StrEnum, IntEnum

from mage_ai.api.operations.constants import OperationType
from mage_ai.data_preparation.models.constants import BlockType, PipelineType


class EntityName(str, Enum):
class EntityName(StrEnum):
ALL = 'ALL'
ALL_EXCEPT_RESERVED = 'ALL_EXCEPT_RESERVED'
AutocompleteItem = 'AutocompleteItem'
Expand Down Expand Up @@ -103,7 +103,7 @@ class EntityName(str, Enum):
]


class BaseEntityType(str, Enum):
class BaseEntityType(StrEnum):
pass


Expand Down Expand Up @@ -131,7 +131,7 @@ class PipelineEntityType(BaseEntityType):
STREAMING = PipelineType.STREAMING.value


class PermissionAccess(int, Enum):
class PermissionAccess(IntEnum):
OWNER = 1
ADMIN = 2
# Editor: list, detail, create, update, delete
Expand Down Expand Up @@ -171,7 +171,7 @@ class PermissionAccess(int, Enum):
DISABLE_UNLESS_CONDITIONS = 1073741824


class PermissionCondition(str, Enum):
class PermissionCondition(StrEnum):
HAS_NOTEBOOK_EDIT_ACCESS = 'HAS_NOTEBOOK_EDIT_ACCESS'
HAS_PIPELINE_EDIT_ACCESS = 'HAS_PIPELINE_EDIT_ACCESS'
USER_OWNS_ENTITY = 'USER_OWNS_ENTITY'
Expand Down
4 changes: 2 additions & 2 deletions mage_ai/cache/constants.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from enum import Enum
from enum import StrEnum

CACHE_KEY_BLOCKS_TO_PIPELINE_MAPPING = 'blocks_to_pipeline_mapping'
CACHE_KEY_BLOCK_ACTION_OBJECTS_MAPPING = 'block_action_objects_mapping'
Expand All @@ -11,5 +11,5 @@
MAGE_CACHE_DIRECTORY_ENVIRONMENT_VARIABLE_NAME = 'MAGE_CACHE_DIRECTORY'


class CacheItemType(str, Enum):
class CacheItemType(StrEnum):
DBT = 'dbt'
4 changes: 2 additions & 2 deletions mage_ai/cache/dbt/constants.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from enum import Enum
from enum import StrEnum

from mage_ai.data_preparation.models.constants import (
BLOCK_TYPE_DIRECTORY_NAME,
Expand All @@ -19,7 +19,7 @@
PROJECT_FILENAMES = [PROJECT_FILENAME, 'dbt_project.yaml']


class FileType(str, Enum):
class FileType(StrEnum):
MODEL = 'model'
PROFILES = 'profiles'
PROJECT = 'project'
Expand Down
4 changes: 2 additions & 2 deletions mage_ai/cluster_manager/constants.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from enum import Enum
from enum import StrEnum

# ECS environment variables
ECS_CLUSTER_NAME = 'ECS_CLUSTER_NAME'
Expand Down Expand Up @@ -29,7 +29,7 @@
NODE_PORT_SERVICE_TYPE = 'NodePort'


class ClusterType(str, Enum):
class ClusterType(StrEnum):
EMR = 'emr'
ECS = 'ecs'
CLOUD_RUN = 'cloud_run'
Expand Down
26 changes: 13 additions & 13 deletions mage_ai/command_center/constants.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
from enum import Enum
from enum import StrEnum

SETTINGS_FILENAME = '.command_center.yaml'


class ItemTagEnum(str, Enum):
class ItemTagEnum(StrEnum):
PINNED = 'pinned'
RECENT = 'recent'


class ItemType(str, Enum):
class ItemType(StrEnum):
ACTION = 'action' # Cast spell
CREATE = 'create' # Conjure
DELETE = 'delete'
Expand All @@ -22,7 +22,7 @@ class ItemType(str, Enum):
UPDATE = 'update'


class ObjectType(str, Enum):
class ObjectType(StrEnum):
APPLICATION = 'application'
APPLICATION_EXPANSION = 'application_expansion'
AUTHENTICATION = 'authentication'
Expand All @@ -42,11 +42,11 @@ class ObjectType(str, Enum):
VERSION_CONTROL_FILE = 'version_control_file'


class ModeType(str, Enum):
class ModeType(StrEnum):
VERSION_CONTROL = 'version_control'


class FileExtension(str, Enum):
class FileExtension(StrEnum):
CSV = 'csv'
JSON = 'json'
MD = 'md'
Expand All @@ -59,7 +59,7 @@ class FileExtension(str, Enum):
YML = 'yml'


class ButtonActionType(str, Enum):
class ButtonActionType(StrEnum):
ADD_APPLICATION = 'add_application'
CLOSE_APPLICATION = 'close_application' # Go back out of the current application.
CLOSE_COMMAND_CENTER = 'close_command_center'
Expand All @@ -70,7 +70,7 @@ class ButtonActionType(str, Enum):
SELECT_ITEM_FROM_REQUEST = 'select_item_from_request'


class InteractionType(str, Enum):
class InteractionType(StrEnum):
CLICK = 'click'
CLOSE_APPLICATION = ButtonActionType.CLOSE_APPLICATION.value
CLOSE_COMMAND_CENTER = ButtonActionType.CLOSE_COMMAND_CENTER.value
Expand All @@ -81,30 +81,30 @@ class InteractionType(str, Enum):
SELECT_ITEM = 'select_item'


class ApplicationType(str, Enum):
class ApplicationType(StrEnum):
DETAIL = 'detail'
DETAIL_LIST = 'detail_list'
EXPANSION = 'expansion'
FORM = 'form'
LIST = 'list'


class ValidationType(str, Enum):
class ValidationType(StrEnum):
CONFIRMATION = 'confirmation'
CUSTOM_VALIDATION_PARSERS = 'custom_validation_parsers'


class RenderLocationType(str, Enum):
class RenderLocationType(StrEnum):
ITEMS_CONTAINER_AFTER = 'items_container_after'


class ApplicationExpansionUUID(str, Enum):
class ApplicationExpansionUUID(StrEnum):
ArcaneLibrary = 'ArcaneLibrary'
PortalTerminal = 'PortalTerminal'
VersionControlFileDiffs = 'VersionControlFileDiffs'


class ApplicationExpansionStatus(str, Enum):
class ApplicationExpansionStatus(StrEnum):
ACTIVE = 'ACTIVE'
CLOSED = 'CLOSED'
INACTIVE = 'INACTIVE'
Expand Down
4 changes: 2 additions & 2 deletions mage_ai/data_cleaner/column_types/constants.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from enum import Enum
from enum import StrEnum


class ColumnType(str, Enum):
class ColumnType(StrEnum):
CATEGORY = 'category'
CATEGORY_HIGH_CARDINALITY = 'category_high_cardinality'
DATETIME = 'datetime'
Expand Down
12 changes: 6 additions & 6 deletions mage_ai/data_cleaner/transformer_actions/constants.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from mage_ai.data_cleaner.column_types.constants import ColumnType
from enum import Enum
from enum import StrEnum
import pandas as pd
import numpy as np
import re
Expand Down Expand Up @@ -34,7 +34,7 @@
}


class ActionType(str, Enum):
class ActionType(StrEnum):
ADD = 'add'
AVERAGE = 'average'
CLEAN_COLUMN_NAME = 'clean_column_name'
Expand Down Expand Up @@ -73,18 +73,18 @@ class ActionType(str, Enum):
STANDARDIZE = 'standardize'


class Axis(str, Enum):
class Axis(StrEnum):
COLUMN = 'column'
ROW = 'row'


class VariableType(str, Enum):
class VariableType(StrEnum):
FEATURE = 'feature'
FEATURE_SET = 'feature_set'
FEATURE_SET_VERSION = 'feature_set_version'


class Operator(str, Enum):
class Operator(StrEnum):
CONTAINS = 'contains'
NOT_CONTAINS = 'not contains'
EQUALS = '=='
Expand All @@ -95,7 +95,7 @@ class Operator(str, Enum):
LESS_THAN_OR_EQUALS = '<='


class ImputationStrategy(str, Enum):
class ImputationStrategy(StrEnum):
AVERAGE = 'average'
COLUMN = 'column'
CONSTANT = 'constant'
Expand Down
4 changes: 2 additions & 2 deletions mage_ai/data_preparation/logging/__init__.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
from dataclasses import dataclass, field
from enum import Enum
from enum import StrEnum
from typing import Dict

from mage_ai.shared.config import BaseConfig
from mage_ai.shared.logger import LoggingLevel


class LoggerType(str, Enum):
class LoggerType(StrEnum):
DEFAULT = 'file'
S3 = 's3'
GCS = 'gcs'
Expand Down
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from enum import Enum
from enum import StrEnum

BLOCK_CATALOG_FILENAME = 'catalog.json'
REPLICATION_METHOD_INCREMENTAL = 'INCREMENTAL'
Expand Down Expand Up @@ -81,6 +81,6 @@
VARIABLE_BOOKMARK_VALUES_KEY = '__bookmark_values__'


class IngestMode(str, Enum):
class IngestMode(StrEnum):
DISK = 'disk'
MEMORY = 'memory'
6 changes: 3 additions & 3 deletions mage_ai/data_preparation/models/block/dbt/constants.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from enum import Enum
from enum import StrEnum

DBT_DIRECTORY_NAME = 'dbt'

Expand All @@ -13,12 +13,12 @@
]


class Flag(str, Enum):
class Flag(StrEnum):
PROFILES_DIR = 'profiles-dir'
PROJECT_DIR = 'project-dir'


class LogLevel(str, Enum):
class LogLevel(StrEnum):
DEBUG = 'debug'
INFO = 'info'
WARN = 'warn'
Expand Down
Loading

0 comments on commit 7dfa7f7

Please sign in to comment.