Add code

Author: Joe Snell

README.md

@@ -1,3 +1,19 @@
 # Python Extension Framework
 
 This is a framework for building [AWS Lambda Extensions](https://aws.amazon.com/blogs/compute/introducing-aws-lambda-extensions-in-preview/).
+
+## Quickstart
+
+To get started you can use the default `Extension` class, or extend it.
+
+Example:
+
+```python
+from lambda_extension_framework import Extension
+
+def handler(event):
+    print(event)
+
+extension = Extension()
+extension.register([EventType.INVOKE], handler)
+```

lambda_extension_framework/main.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+
+import os
+import sys
+from pathlib import Path
+from enum import Enum
+
+from lambda_extension_framework.utils import get, post
+
+class EventType(Enum):
+    INVOKE='INVOKE'
+    SHUTDOWN='SHUTDOWN'
+
+class Extension:
+    def __init__(self, name=None, runtime_api_endpoint=None):
+        self.handlers = {}
+        self.name = name or Path(__file__).parent.name
+        self.runtime_api_endpoint = runtime_api_endpoint or os.environ['AWS_LAMBDA_RUNTIME_API']
+
+    @staticmethod
+    def handle_signal(signal, frame):
+        sys.exit(0)
+
+    def register(self, events, handler):
+        response = post(
+            url=f'https://{self.runtime_api_endpoint}/2020-01-01/extension/register',
+            data={
+                'events': events
+            },
+            headers={
+                'Lambda-Extension-Name': self.name,
+            }
+        )
+    
+        extension_id = response.headers['Lambda-Extension-Identifier']
+        return self.process_events(extension_id, handler)
+
+    def process_events(self, extension_id, handler):
+        while True:
+            response = get(
+                url=f"https://{self.runtime_api_endpoint}/2020-01-01/extension/event/next",
+                headers={
+                    'Lambda-Extension-Identifier': extension_id
+                },
+                timeout=None
+            )
+
+            if response['eventType'] == EventType.SHUTDOWN:
+                sys.exit(0)
+            else:
+                handler(response)

lambda_extension_framework/utils.py

@@ -0,0 +1,26 @@
+import json
+from urllib import request
+
+def get(url, headers=None):
+    kwargs = {
+        'url': url,
+        'headers': headers or {}
+    }
+
+    req = request.Request(**kwargs)                    
+    response = request.urlopen(req)
+    return response.read().decode('utf8')
+
+def post(url, headers=None, data=None):
+    kwargs = {
+        'url': url,
+        'headers': headers or {}
+    }
+
+    if data:
+        kwargs['data'] = json.dumps(data)
+        kwargs['headers']['Content-Type'] = 'application/json'
+
+    req = request.Request(**kwargs)                    
+    response = request.urlopen(req)
+    return response.read().decode('utf8')

pyproject.toml

@@ -1,9 +1,20 @@
 [tool.poetry]
-name = "python-extension-framework"
+name = "lambda-extension-framework"
 version = "0.1.0"
-description = ""
+description = "Framework for creating AWS Lambda Extensions"
 authors = ["Joe Snell <joepsnell@gmail.com>"]
 license = "MIT"
+readme = "README.md"
+repository = "https://github.com/lambda-extensions/python-extension-framework"
+documentation = "https://github.com/lambda-extensions/python-extension-framework"
+homepage = "https://github.com/lambda-extensions/python-extension-framework"
+keywords = [
+    "AWS",
+    "AWS Lambda",
+    "Lambda Extensions",
+    "Framework",
+    "Extension Framework"
+]
 
 [tool.poetry.dependencies]
 python = "^3.8"