Repository: tobymao/sqlglot

Issue:
Unable to use Dialect instances as a read/write arguments when transpiling
Although typing indicates that functions like `transpile` accept instances of `Dialect` as `read` and `write` arguments, it doesn't seem to work. I'd be happy to submit a merge request to fix this, but it seems like this issue goes deeper, with `Dialect.get_or_raise` actually being the root problem, as it seems to be used both to access a certain Dialect class, as well as generate Dialect objects. Defining typing for that method would probably highlight this issue well.

Since I need transpiling to work with my custom dialects, I'd be happy to work on this, but I don't want to mess something up before someone maintaining the package confirms that this is indeed an issue and that I'm on the right track.

p.s. awesome package!

Discussion:

Patch:
diff --git a/sqlglot/__init__.py b/sqlglot/__init__.py
index 6861590a..6aefb27c 100644
--- a/sqlglot/__init__.py
+++ b/sqlglot/__init__.py
@@ -34,8 +34,11 @@ from sqlglot.schema import MappingSchema, Schema
 from sqlglot.tokens import Tokenizer, TokenType
 
 if t.TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
+
     T = t.TypeVar("T", bound=Expression)
 
+
 __version__ = "10.6.0"
 
 pretty = False
@@ -45,9 +48,7 @@ schema = MappingSchema()
 """The default schema used by SQLGlot (e.g. in the optimizer)."""
 
 
-def parse(
-    sql: str, read: t.Optional[str | Dialect] = None, **opts
-) -> t.List[t.Optional[Expression]]:
+def parse(sql: str, read: DialectType = None, **opts) -> t.List[t.Optional[Expression]]:
     """
     Parses the given SQL string into a collection of syntax trees, one per parsed SQL statement.
 
@@ -76,7 +77,7 @@ def parse_one(
 @t.overload
 def parse_one(
     sql: str,
-    read: str | Dialect,
+    read: DialectType,
     into: t.Type[T],
     **opts,
 ) -> T:
@@ -96,7 +97,7 @@ def parse_one(
 @t.overload
 def parse_one(
     sql: str,
-    read: str | Dialect,
+    read: DialectType,
     into: t.Union[str, t.Collection[t.Union[str, t.Type[Expression]]]],
     **opts,
 ) -> Expression:
@@ -113,7 +114,7 @@ def parse_one(
 
 def parse_one(
     sql: str,
-    read: t.Optional[str | Dialect] = None,
+    read: DialectType = None,
     into: t.Optional[exp.IntoType] = None,
     **opts,
 ) -> Expression:
@@ -147,8 +148,8 @@ def parse_one(
 
 def transpile(
     sql: str,
-    read: t.Optional[str | Dialect] = None,
-    write: t.Optional[str | Dialect] = None,
+    read: DialectType = None,
+    write: DialectType = None,
     identity: bool = True,
     error_level: t.Optional[ErrorLevel] = None,
     **opts,
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
index 0c2bebab..aed494b8 100644
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -11,6 +11,9 @@ from sqlglot.time import format_time
 from sqlglot.tokens import Tokenizer
 from sqlglot.trie import new_trie
 
+if t.TYPE_CHECKING:
+    DialectType = t.Union[str, Dialect, t.Type[Dialect], None]
+
 
 class Dialects(str, Enum):
     DIALECT = ""
@@ -122,9 +125,15 @@ class Dialect(metaclass=_Dialect):
     def get_or_raise(cls, dialect):
         if not dialect:
             return cls
+        if isinstance(dialect, _Dialect):
+            return dialect
+        if isinstance(dialect, Dialect):
+            return dialect.__class__
+
         result = cls.get(dialect)
         if not result:
             raise ValueError(f"Unknown dialect '{dialect}'")
+
         return result
 
     @classmethod
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 7c1a1167..da06d1fc 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -32,7 +32,7 @@ from sqlglot.helper import (
 from sqlglot.tokens import Token
 
 if t.TYPE_CHECKING:
-    from sqlglot.dialects.dialect import Dialect
+    from sqlglot.dialects.dialect import DialectType
 
     IntoType = t.Union[
         str,
@@ -427,7 +427,7 @@ class Expression(metaclass=_Expression):
     def __repr__(self):
         return self._to_s()
 
-    def sql(self, dialect: Dialect | str | None = None, **opts) -> str:
+    def sql(self, dialect: DialectType = None, **opts) -> str:
         """
         Returns SQL string representation of this tree.
 
@@ -2455,9 +2455,7 @@ class DataType(Expression):
     }
 
     @classmethod
-    def build(
-        cls, dtype: str | DataType.Type, dialect: t.Optional[str | Dialect] = None, **kwargs
-    ) -> DataType:
+    def build(cls, dtype: str | DataType.Type, dialect: DialectType = None, **kwargs) -> DataType:
         from sqlglot import parse_one
 
         if isinstance(dtype, str):
@@ -3476,7 +3474,7 @@ def maybe_parse(
     sql_or_expression: str | Expression,
     *,
     into: t.Optional[IntoType] = None,
-    dialect: t.Optional[str] = None,
+    dialect: DialectType = None,
     prefix: t.Optional[str] = None,
     **opts,
 ) -> Expression:
@@ -4373,7 +4371,7 @@ def expand(expression: Expression, sources: t.Dict[str, Subqueryable], copy=True
     return expression.transform(_expand, copy=copy)
 
 
-def func(name: str, *args, dialect: t.Optional[Dialect | str] = None, **kwargs) -> Func:
+def func(name: str, *args, dialect: DialectType = None, **kwargs) -> Func:
     """
     Returns a Func expression.
 
diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
index 4e7eab84..a39ad8cf 100644
--- a/sqlglot/lineage.py
+++ b/sqlglot/lineage.py
@@ -9,6 +9,9 @@ from sqlglot.optimizer import Scope, build_scope, optimize
 from sqlglot.optimizer.qualify_columns import qualify_columns
 from sqlglot.optimizer.qualify_tables import qualify_tables
 
+if t.TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
+
 
 @dataclass(frozen=True)
 class Node:
@@ -36,7 +39,7 @@ def lineage(
     schema: t.Optional[t.Dict | Schema] = None,
     sources: t.Optional[t.Dict[str, str | exp.Subqueryable]] = None,
     rules: t.Sequence[t.Callable] = (qualify_tables, qualify_columns),
-    dialect: t.Optional[str] = None,
+    dialect: DialectType = None,
 ) -> Node:
     """Build the lineage graph for a column of a SQL query.
 
@@ -126,7 +129,7 @@ class LineageHTML:
     def __init__(
         self,
         node: Node,
-        dialect: t.Optional[str] = None,
+        dialect: DialectType = None,
         imports: bool = True,
         **opts: t.Any,
     ):
diff --git a/sqlglot/schema.py b/sqlglot/schema.py
index f6f38836..f5d9f2b0 100644
--- a/sqlglot/schema.py
+++ b/sqlglot/schema.py
@@ -11,6 +11,7 @@ from sqlglot.trie import in_trie, new_trie
 
 if t.TYPE_CHECKING:
     from sqlglot.dataframe.sql.types import StructType
+    from sqlglot.dialects.dialect import DialectType
 
     ColumnMapping = t.Union[t.Dict, str, StructType, t.List]
 
@@ -153,7 +154,7 @@ class MappingSchema(AbstractMappingSchema[t.Dict[str, str]], Schema):
         self,
         schema: t.Optional[t.Dict] = None,
         visible: t.Optional[t.Dict] = None,
-        dialect: t.Optional[str] = None,
+        dialect: DialectType = None,
     ) -> None:
         self.dialect = dialect
         self.visible = visible or {}


Test patch:
diff --git a/tests/dialects/test_dialect.py b/tests/dialects/test_dialect.py
index 5a13655c..446247c3 100644
--- a/tests/dialects/test_dialect.py
+++ b/tests/dialects/test_dialect.py
@@ -1,6 +1,7 @@
 import unittest
 
 from sqlglot import Dialect, Dialects, ErrorLevel, UnsupportedError, parse_one
+from sqlglot.dialects import Hive
 
 
 class Validator(unittest.TestCase):
@@ -67,6 +68,11 @@ class TestDialect(Validator):
             self.assertIsNotNone(Dialect.get_or_raise(dialect))
             self.assertIsNotNone(Dialect[dialect.value])
 
+    def test_get_or_raise(self):
+        self.assertEqual(Dialect.get_or_raise(Hive), Hive)
+        self.assertEqual(Dialect.get_or_raise(Hive()), Hive)
+        self.assertEqual(Dialect.get_or_raise("hive"), Hive)
+
     def test_cast(self):
         self.validate_all(
             "CAST(a AS TEXT)",