From 539138de40e59cebc7b0b926ce515c6fa341c837 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Paul=20Pr=C3=BCnster?=
 <37049593+PaulPruenster@users.noreply.github.com>
Date: Tue, 2 Jul 2024 10:49:52 +0200
Subject: [PATCH] Missing-observation-autocomplete (#231)

* add types

* return type
---
 rlbench/demo.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/rlbench/demo.py b/rlbench/demo.py
index bf3df6a0d..18d32f0fc 100644
--- a/rlbench/demo.py
+++ b/rlbench/demo.py
@@ -1,9 +1,11 @@
 import numpy as np
+from typing import List
+from rlbench.backend.observation import Observation
 
 
 class Demo(object):
 
-    def __init__(self, observations, random_seed=None, num_reset_attempts = None):
+    def __init__(self, observations: List[Observation], random_seed=None, num_reset_attempts=None):
         self._observations = observations
         self.random_seed = random_seed
         self.num_reset_attempts = num_reset_attempts
@@ -11,7 +13,7 @@ def __init__(self, observations, random_seed=None, num_reset_attempts = None):
     def __len__(self):
         return len(self._observations)
 
-    def __getitem__(self, i):
+    def __getitem__(self, i) -> Observation:
         return self._observations[i]
 
     def restore_state(self):