Search code examples
pythonamazon-web-servicesclassaws-lambdaenvironment-variables

How to set up different sets of class methods in based on an environment variable?


I am developing a toy crypto project that involves training keras models and then deploying them in AWS lambda. I test the models locally and deploy them as lambda functions in AWS. Some classes are used both locally and during lambda invocation, but they have wildly different internal working. For instance, a class Asset implements current_price method, but it calls API in one case and local db in another.

I thought of using environment variables to distinguish two environments. How do I instantiate the same class with a different set of methods or different method implementations based on an environment variable? Is there a better approach to this problem? I'm writing in python.

P.S. Based on the comments and answers, I ended up implementing the following solution:

import os

class Base:
  def __init__(self, param_base = 'Base'):
    self.param_base = param_base

  def method_base(self):
    print(self.param_base)

class Asset_local(Base):
  def __init__(self, param_local = 'Loc'):
    super().__init__()
    self.param = param_local

  def method_local(self):
    print(self.param)

class Asset_lambda(Base):
  def __init__(self, param_lambda = 'Lambda'):
    super().__init__()
    self.param = param_lambda

  def method_lambda(self):
    print(self.param)

class Asset():
  def __new__(cls):
    env_var = os.environ['MY_VAR']
    if env_var == 'Lambda':
        return Asset_lambda()
    if env_var == 'Local':
        return Asset_local()

This allows me to keep the common parts of the code in the parent class, implement the individual functions in the separate classes and keep all of that hidden from the code that calls the class.


Solution

  • You can define methods conditionally...

    class MyClass:
        if some_condition:
            def some_method(self):
                ...
        else:
            def another_method(self):
                ...
    

    But this is kind of awkward. It's usually better for a class to always have the same methods and their behavior can be controlled by parameters, composition, or some other means.

    class Myclass:
        def __init__(self, execution_mode: Literal['lambda', 'local']):
            assert execution_mode in ('lambda', 'local'), "Invalid execution mode"
            self.execution_mode = execution_mode
    
        def _some_method_lambda(self):
            ... # implementation when running in lambda
    
        def _some_method_local(self):
            ... # implementation when running locally
    
        def some_method(self, *args, **kwargs):
            if self.execution_mode == 'local':
                self._some_method_local(*args, **kwargs)
            elif self.execution_mode == 'lambda':
                self._some_method_lambda(*args, **kwargs)
    ...
    
    if os.environ.get('EXECUTION_ENVIRONMENT') == 'lambda':
        my_instance = MyClass(execution_mode='lambda')
    else:
        my_instance = MyClass(execution_mode='local')
    

    Or something like this. There's a lot of ways you can slice this.