Keep Optional and union in valid_static_types

I'm keeping them here so as to avoid a potential bug, it's used in type
checking and it makes sense to keep it, I've added a second type for
Union using UnionType.

Consistency with the placement of None for optional arguments has also
been improved.

Signed-off-by: Chihurumnaya Ibiam <ibiamchihurumnaya@gmail.com>

Author: chimosky

osprey_worker/src/osprey/engine/ast/grammar.py

@@ -350,7 +350,7 @@ class Assign(Statement, IsConstant, IsExtractable):
 
     target: Name
     value: Expression
-    annotation: None | 'Annotation' | 'AnnotationWithVariants' = None
+    annotation: 'Annotation' | 'AnnotationWithVariants' | None = None
 
     @cached_property
     def should_extract(self) -> bool:
@@ -411,7 +411,7 @@ class Call(Expression, Statement):
     func: Name | 'Attribute'
     arguments: Sequence['Keyword']
 
-    def find_argument(self, name: str) -> None | 'Keyword':
+    def find_argument(self, name: str) -> 'Keyword' | None:
         for argument in self.arguments:
             if argument.name == name:
                 return argument

osprey_worker/src/osprey/engine/ast_validator/validators/validate_static_types.py

@@ -2,7 +2,7 @@
 from dataclasses import dataclass, replace
 from functools import lru_cache
 from types import UnionType
-from typing import TYPE_CHECKING, Any, Type, cast
+from typing import TYPE_CHECKING, Any, Optional, Type, Union, cast
 
 from osprey.engine.ast import grammar
 from osprey.engine.ast.error_utils import SpanWithHint
@@ -452,9 +452,9 @@ def _is_optional_type(self, t: type) -> bool:
 
         # We allow Entities to typecheck as optional in order to avoid having to update the type system
         # to support Optional[Entity[str]]
-        if origin is None or origin is EntityT:
+        if origin is Optional or origin is EntityT:
             return True
-        elif origin is UnionType:
+        elif origin is Union or origin is UnionType:
             return type(None) in t.__args__  # type: ignore[attr-defined]
 
         return False
@@ -613,7 +613,7 @@ def _validate_format_string(self, format_string: grammar.FormatString) -> type:
             accepted_types_str = ', '.join(to_display_str(arg) for arg in accepted_types)
             self._check_compatible_type(
                 type_t=name_type,
-                accepted_by_t=cast(type, accepted_types),
+                accepted_by_t=cast(type, Union[accepted_types]),
                 message='unsupported type for f-string substitution',
                 node=name,
                 hint=f'has type {name_type_str}, expected one of {accepted_types_str}',